Skip to content

Commit 9ce222e

Browse files
SimonRastikianSimonRastikian
authored
refactor: simplifying storage by giving it to the client directly - #3851 (#4227)
1 parent d111b40 commit 9ce222e

6 files changed

Lines changed: 53 additions & 47 deletions

File tree

‎crates/node/src/assets/cleanup.rs‎

Lines changed: 10 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -132,13 +132,12 @@ mod tests {
132132
use crate::assets::cleanup::EpochData;
133133
use crate::assets::cleanup::{delete_stale_triples_and_presignatures, get_epoch_data};
134134
use crate::assets::test_utils;
135-
use crate::assets::test_utils::TestContext;
136-
use crate::assets::test_utils::get_participant_ids;
137-
use crate::assets::test_utils::make_triple;
138-
use crate::assets::test_utils::random_verifying_key;
139-
use crate::assets::test_utils::triple_v2_key;
135+
use crate::assets::test_utils::{
136+
TestContext, get_participant_ids, make_triple, random_verifying_key, triple_v2_key,
137+
};
140138
use crate::db::EPOCH_ID_KEY;
141139
use crate::db::{DBCol, SecretDB};
140+
use crate::network::testing::new_test_client;
142141
use crate::primitives::UniqueId;
143142
use crate::providers::ecdsa::triple::TripleStorage;
144143
use mpc_primitives::domain::DomainId;
@@ -241,17 +240,13 @@ mod tests {
241240
let (mut start_data, my_participant_id, reconstruction_threshold) =
242241
test_utils::gen_four_participants();
243242
let all_participants = get_participant_ids(start_data.clone());
244-
let alive_participants = Arc::new(Mutex::new(all_participants.clone()));
245243
let dir = tempfile::tempdir().unwrap();
246244
let db = SecretDB::new(dir.path(), [1; 16]).unwrap();
245+
let client = new_test_client(all_participants.clone(), my_participant_id);
247246
let triple_store = TripleStorage::new(
248247
FakeClock::default().clock(),
249248
db.clone(),
250-
my_participant_id,
251-
{
252-
let alive = alive_participants.clone();
253-
Arc::new(move || alive.lock().unwrap().clone())
254-
},
249+
client,
255250
reconstruction_threshold,
256251
)
257252
.unwrap();
@@ -310,17 +305,13 @@ mod tests {
310305
.take(all_participants.len() - 1)
311306
.copied()
312307
.collect();
313-
let alive_participants = Arc::new(Mutex::new(all_participants.clone()));
314308
let dir = tempfile::tempdir().unwrap();
315309
let db = SecretDB::new(dir.path(), [1; 16]).unwrap();
310+
let client = new_test_client(all_participants.clone(), my_participant_id);
316311
let triple_store = TripleStorage::new(
317312
FakeClock::default().clock(),
318313
db.clone(),
319-
my_participant_id,
320-
{
321-
let alive = alive_participants.clone();
322-
Arc::new(move || alive.lock().unwrap().clone())
323-
},
314+
client,
324315
reconstruction_threshold,
325316
)
326317
.unwrap();
@@ -369,17 +360,13 @@ mod tests {
369360
.iter()
370361
.find(|p| **p != my_participant_id)
371362
.unwrap();
372-
let alive_participants = Arc::new(Mutex::new(all_participants.clone()));
373363
let dir = tempfile::tempdir().unwrap();
374364
let db = SecretDB::new(dir.path(), [1; 16]).unwrap();
365+
let client = new_test_client(all_participants.clone(), my_participant_id);
375366
let triple_store = TripleStorage::new(
376367
FakeClock::default().clock(),
377368
db.clone(),
378-
my_participant_id,
379-
{
380-
let alive = alive_participants.clone();
381-
Arc::new(move || alive.lock().unwrap().clone())
382-
},
369+
client,
383370
reconstruction_threshold,
384371
)
385372
.unwrap();

‎crates/node/src/network.rs‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ pub mod handshake;
55
pub mod indexer_heights;
66

77
use crate::metrics::networking_metrics;
8+
use crate::network::indexer_heights::IndexerHeightTracker;
89
use crate::primitives::{
910
ChannelId, IndexerHeightMessage, MpcMessage, MpcMessageKind, MpcPeerMessage, MpcStartMessage,
1011
MpcTaskId, ParticipantId, PeerMessage, UniqueId,
@@ -13,7 +14,6 @@ use crate::requests::queue::NetworkAPIForRequests;
1314
use crate::tracking::{self, AutoAbortTask};
1415
use anyhow::Context as _;
1516
use conn::{ConnectionVersion, NodeConnectivityInterface};
16-
use indexer_heights::IndexerHeightTracker;
1717
use lru::LruCache;
1818
use rand::prelude::IteratorRandom;
1919
use std::collections::hash_map::Entry;
@@ -866,6 +866,7 @@ impl NetworkTaskChannel {
866866
#[cfg(test)]
867867
pub mod testing {
868868
use super::conn::{ConnectionVersion, NodeConnectivityInterface};
869+
use super::indexer_heights::IndexerHeightTracker;
869870
use super::{
870871
ChannelId, MeshNetworkTransportSender, NetworkTaskChannel, NetworkTaskChannelSender,
871872
};
@@ -1002,6 +1003,29 @@ pub mod testing {
10021003
transports
10031004
}
10041005

1006+
/// Synchronous [`MeshNetworkClient`] for unit tests. All participants are reported alive.
1007+
pub fn new_test_client(
1008+
participants: Vec<ParticipantId>,
1009+
my_participant_id: ParticipantId,
1010+
) -> Arc<super::MeshNetworkClient> {
1011+
let transport = Arc::new(TestMeshTransportSender {
1012+
transport: Arc::new(TestMeshTransport {
1013+
participant_ids: participants.clone(),
1014+
senders: HashMap::new(),
1015+
}),
1016+
my_participant_id,
1017+
});
1018+
let channels = Arc::new(std::sync::Mutex::new(
1019+
super::NetworkTaskChannelManager::new(),
1020+
));
1021+
let indexer_heights = Arc::new(IndexerHeightTracker::new(&participants));
1022+
Arc::new(super::MeshNetworkClient::new(
1023+
transport,
1024+
channels,
1025+
indexer_heights,
1026+
))
1027+
}
1028+
10051029
/// Builds a channel over the given participant set, returning the raw inbound sender so
10061030
/// tests can inject arbitrary [`MpcPeerMessage`]s, including ones from outside the set.
10071031
pub fn new_task_channel_for_test(

‎crates/node/src/providers/ecdsa.rs‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,7 @@ use crate::db::SecretDB;
1616
use crate::metrics::tokio_task_metrics::ECDSA_TASK_MONITORS;
1717
use crate::network::{MeshNetworkClient, NetworkTaskChannel};
1818
use crate::primitives::{MpcTaskId, ParticipantId, UniqueId};
19-
use crate::providers::DomainKeyshare;
20-
use crate::providers::SignatureProvider;
21-
use crate::providers::ecdsa_common;
19+
use crate::providers::{DomainKeyshare, SignatureProvider, ecdsa_common};
2220
use crate::storage::SignRequestStorage;
2321
use crate::tracking;
2422
use mpc_node_config::ConfigFile;
@@ -60,7 +58,7 @@ impl EcdsaSignatureProvider {
6058
sign_request_store: Arc<SignRequestStorage>,
6159
keyshares: HashMap<DomainId, DomainKeyshare<Secp256K1Sha256>>,
6260
) -> anyhow::Result<Self> {
63-
let keyshares = ecdsa_common::build_keyshares(&clock, &db, &client, keyshares)?;
61+
let keyshares = ecdsa_common::build_keyshares(&clock, &db, client.clone(), keyshares)?;
6462

6563
// cait-sith triple generation runs with exactly `t` parties, so keep one store per distinct reconstruction threshold.
6664
let mut triple_stores = HashMap::new();
@@ -71,8 +69,7 @@ impl EcdsaSignatureProvider {
7169
Arc::new(TripleStorage::new(
7270
clock.clone(),
7371
db.clone(),
74-
client.my_participant_id(),
75-
ecdsa_common::active_participants_query(&client),
72+
client.clone(),
7673
t,
7774
)?),
7875
);

‎crates/node/src/providers/ecdsa/triple.rs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ use crate::primitives::{ParticipantId, UniqueId};
1010
use crate::protocol::NamedProtocol;
1111
use crate::providers::HasParticipants;
1212
use crate::providers::ecdsa::{EcdsaSignatureProvider, EcdsaTaskId};
13+
use crate::providers::ecdsa_common::active_participants_query;
1314
use crate::tracking::AutoAbortTaskCollection;
1415
use mpc_node_config::TripleConfig;
1516
use mpc_primitives::ReconstructionThreshold;
@@ -55,18 +56,17 @@ impl TripleStorage {
5556
pub fn new(
5657
clock: Clock,
5758
db: Arc<SecretDB>,
58-
my_participant_id: ParticipantId,
59-
alive_participant_ids_query: Arc<dyn Fn() -> Vec<ParticipantId> + Send + Sync>,
59+
client: Arc<MeshNetworkClient>,
6060
reconstruction_threshold: ReconstructionThreshold,
6161
) -> anyhow::Result<Self> {
6262
Ok(Self(DistributedAssetStorage::<PairedTriple>::new(
6363
clock,
6464
db,
6565
DBCol::TripleV2,
6666
reconstruction_threshold.inner().to_be_bytes().to_vec(),
67-
my_participant_id,
67+
client.my_participant_id(),
6868
|participants, pair| pair.is_subset_of_active_participants(participants),
69-
alive_participant_ids_query,
69+
active_participants_query(client),
7070
)?))
7171
}
7272
}
@@ -385,7 +385,7 @@ mod tests {
385385
use crate::assets::test_utils::{make_triple, triple_v2_key};
386386
use crate::db::{DBCol, SecretDB};
387387
use crate::network::computation::MpcLeaderCentricComputation;
388-
use crate::network::testing::run_test_clients;
388+
use crate::network::testing::{new_test_client, run_test_clients};
389389
use crate::network::{MeshNetworkClient, NetworkTaskChannel};
390390
use crate::primitives::{MpcTaskId, ParticipantId, UniqueId};
391391
use crate::providers::ecdsa::EcdsaTaskId;
@@ -569,13 +569,13 @@ mod tests {
569569
db: Arc<SecretDB>,
570570
my_participant_id: ParticipantId,
571571
reconstruction_threshold: ReconstructionThreshold,
572-
alive: Vec<ParticipantId>,
572+
participants: Vec<ParticipantId>,
573573
) -> TripleStorage {
574+
let client = new_test_client(participants, my_participant_id);
574575
TripleStorage::new(
575576
near_time::FakeClock::default().clock(),
576577
db,
577-
my_participant_id,
578-
Arc::new(move || alive.clone()),
578+
client,
579579
reconstruction_threshold,
580580
)
581581
.unwrap()

‎crates/node/src/providers/ecdsa_common.rs‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ where
4848
pub fn new(
4949
clock: Clock,
5050
db: Arc<SecretDB>,
51-
client: &Arc<MeshNetworkClient>,
51+
client: Arc<MeshNetworkClient>,
5252
domain_id: DomainId,
5353
) -> anyhow::Result<Self> {
5454
Ok(Self(DistributedAssetStorage::<
@@ -94,17 +94,16 @@ where
9494

9595
/// The "are all these participants still alive?" query both the presignature and triple stores use.
9696
pub fn active_participants_query(
97-
client: &Arc<MeshNetworkClient>,
97+
client: Arc<MeshNetworkClient>,
9898
) -> Arc<dyn Fn() -> Vec<ParticipantId> + Send + Sync> {
99-
let network_client = client.clone();
100-
Arc::new(move || network_client.all_alive_participant_ids())
99+
Arc::new(move || client.all_alive_participant_ids())
101100
}
102101

103102
/// Attaches a freshly-created presignature store to each domain's [`DomainKeyshare`].
104103
pub fn build_keyshares<P>(
105104
clock: &Clock,
106105
db: &Arc<SecretDB>,
107-
client: &Arc<MeshNetworkClient>,
106+
client: Arc<MeshNetworkClient>,
108107
keyshares: HashMap<DomainId, DomainKeyshare<Secp256K1Sha256>>,
109108
) -> anyhow::Result<HashMap<DomainId, EcdsaKeyshare<P>>>
110109
where
@@ -115,7 +114,7 @@ where
115114
let presignature_store = Arc::new(PresignatureStorage::new(
116115
clock.clone(),
117116
db.clone(),
118-
client,
117+
client.clone(),
119118
domain_id,
120119
)?);
121120
result.insert(
@@ -197,8 +196,7 @@ mod tests {
197196

198197
// When
199198
let keyshares =
200-
build_keyshares::<Vec<u8>>(&Clock::real(), &db, &client, keyshares)
201-
.unwrap();
199+
build_keyshares::<Vec<u8>>(&Clock::real(), &db, client, keyshares).unwrap();
202200

203201
// Then each domain keeps the threshold it was configured with
204202
assert_eq!(keyshares[&low].reconstruction_threshold, low_threshold);

‎crates/node/src/providers/robust_ecdsa.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ impl RobustEcdsaSignatureProvider {
6060
sign_request_store: Arc<SignRequestStorage>,
6161
keyshares: HashMap<DomainId, DomainKeyshare<Secp256K1Sha256>>,
6262
) -> anyhow::Result<Self> {
63-
let keyshares = ecdsa_common::build_keyshares(&clock, &db, &client, keyshares)?;
63+
let keyshares = ecdsa_common::build_keyshares(&clock, &db, client.clone(), keyshares)?;
6464

6565
Ok(Self {
6666
config,

0 commit comments

Comments
 (0)