Skip to content

Commit dcb24d5

Browse files
ManuelBilbaoManuelBilbao
committed
Optimize initial setup reducing requests
1 parent 16ca00d commit dcb24d5

File tree

2 files changed

+141
-100
lines changed

2 files changed

+141
-100
lines changed

‎crates/common/src/config/signer.rs‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,6 @@ pub struct DirkHostConfig {
4343
pub accounts: Vec<String>,
4444
}
4545

46-
impl DirkHostConfig {
47-
pub fn name(&self) -> Option<String> {
48-
self.server_name.clone().or_else(|| self.url.host_str().map(String::from))
49-
}
50-
}
51-
5246
#[derive(Debug, Serialize, Deserialize, Clone)]
5347
#[serde(rename_all = "snake_case")]
5448
pub enum SignerType {

‎crates/signer/src/manager/dirk.rs‎

Lines changed: 141 additions & 94 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
1-
use std::{collections::HashMap, io::Write, path::PathBuf};
1+
use std::{
2+
collections::{HashMap, HashSet},
3+
io::Write,
4+
path::PathBuf,
5+
};
26

37
use alloy::{
48
hex, rpc::types::beacon::constants::BLS_SIGNATURE_BYTES_LEN, transports::http::reqwest::Url,
@@ -101,11 +105,6 @@ impl DirkManager {
101105
let mut consensus_accounts = HashMap::new();
102106

103107
for host in config.hosts {
104-
let Some(host_name) = host.name() else {
105-
warn!("Host name not found for server {}", host.url);
106-
continue;
107-
};
108-
109108
let channel = match connect(&host, &certs).await {
110109
Ok(channel) => channel,
111110
Err(e) => {
@@ -116,93 +115,44 @@ impl DirkManager {
116115

117116
connections.insert(host.url.clone(), channel.clone());
118117

119-
// TODO: Improve to minimize requests
120-
for account_name in host.accounts {
121-
let Ok((wallet, name)) = decompose_name(&account_name) else {
122-
warn!("Invalid account name {account_name}");
123-
continue;
124-
};
125-
126-
let response = match ListerClient::new(channel.clone())
127-
.list_accounts(ListAccountsRequest { paths: vec![account_name.clone()] })
128-
.await
129-
{
130-
Ok(res) => res,
131-
Err(e) => {
132-
warn!("Failed to get account {account_name}: {e}");
133-
continue;
134-
}
135-
};
118+
let wallets: HashSet<String> = host
119+
.accounts
120+
.iter()
121+
.map(|account| decompose_name(account).unwrap_or_default().0)
122+
.collect();
136123

137-
if response.get_ref().state() != ResponseState::Succeeded {
138-
warn!("Failed to get account {account_name}");
124+
let accounts_response = match ListerClient::new(channel.clone())
125+
.list_accounts(ListAccountsRequest { paths: wallets.into_iter().collect() })
126+
.await
127+
{
128+
Ok(res) => res,
129+
Err(e) => {
130+
warn!("Failed to list accounts in server {}: {e}", host.url);
139131
continue;
140132
}
133+
};
141134

142-
if let Some(account) = response.get_ref().accounts.get(0) {
143-
// The account is Simple
144-
match BlsPublicKey::try_from(account.public_key.as_slice()) {
145-
Ok(public_key) => {
146-
consensus_accounts.insert(
147-
public_key,
148-
Account::Simple(SimpleAccount {
149-
public_key,
150-
server: host.url.clone(),
151-
wallet: wallet.to_string(),
152-
name: name.to_string(),
153-
}),
154-
);
155-
}
156-
Err(_) => {
157-
warn!("Failed to parse public key for account {account_name}");
158-
continue;
159-
}
160-
}
161-
} else if let Some(account) = response.get_ref().distributed_accounts.get(0) {
162-
// The account is Distributed
163-
let Ok(public_key) =
164-
BlsPublicKey::try_from(account.composite_public_key.as_slice())
165-
else {
166-
warn!("Failed to parse composite public key for account {account_name}");
167-
continue;
168-
};
169-
170-
let Some(&Endpoint { id: participant_id, .. }) = account
171-
.participants
172-
.iter()
173-
.find(|participant| participant.name == host_name)
174-
else {
175-
warn!(
176-
"Host {host_name} not found as participant for account {account_name}"
177-
);
178-
continue;
179-
};
180-
181-
match consensus_accounts.get_mut(&public_key) {
182-
Some(Account::Distributed(DistributedAccount { participants, .. })) => {
183-
participants.insert(participant_id as u32, host.url.clone());
184-
}
185-
Some(Account::Simple(_)) => {
186-
bail!("Distributed public key already exists for simple account");
187-
}
188-
None => {
189-
let mut participants =
190-
HashMap::with_capacity(account.participants.len());
191-
participants.insert(participant_id as u32, host.url.clone());
192-
consensus_accounts.insert(
193-
public_key,
194-
Account::Distributed(DistributedAccount {
195-
composite_public_key: public_key,
196-
participants,
197-
threshold: account.signing_threshold,
198-
wallet: wallet.to_string(),
199-
name: name.to_string(),
200-
}),
201-
);
202-
}
203-
}
204-
} else {
205-
warn!("Account {account_name} not found in server {}", host.url);
135+
if accounts_response.get_ref().state() != ResponseState::Succeeded {
136+
warn!("Failed to list accounts in server {}", host.url);
137+
continue;
138+
}
139+
140+
let accounts_response = accounts_response.into_inner();
141+
load_simple_accounts(accounts_response.accounts, &host, &mut consensus_accounts);
142+
load_distributed_accounts(
143+
accounts_response.distributed_accounts,
144+
&host,
145+
&mut consensus_accounts,
146+
)
147+
.map_err(|error| warn!("{error}"))
148+
.ok();
149+
150+
for account in host.accounts {
151+
if !consensus_accounts
152+
.values()
153+
.any(|account| account.full_name() == account.full_name())
154+
{
155+
warn!("Account {account} not found in server {}", host.url);
206156
}
207157
}
208158
}
@@ -580,13 +530,110 @@ async fn connect(server: &DirkHostConfig, certs: &CertConfig) -> eyre::Result<Ch
580530
.map_err(eyre::Error::from)
581531
}
582532

583-
fn decompose_name(full_name: &str) -> eyre::Result<(&str, &str)> {
584-
full_name.split_once('/').ok_or_else(|| eyre::eyre!("Invalid account name"))
533+
fn decompose_name(full_name: &str) -> eyre::Result<(String, String)> {
534+
full_name
535+
.split_once('/')
536+
.map(|(wallet, name)| (wallet.to_string(), name.to_string()))
537+
.ok_or_else(|| eyre::eyre!("Invalid account name"))
538+
}
539+
540+
fn load_simple_accounts(
541+
accounts: Vec<crate::proto::v1::Account>,
542+
host: &DirkHostConfig,
543+
consensus_accounts: &mut HashMap<BlsPublicKey, Account>,
544+
) {
545+
for account in accounts {
546+
if !host.accounts.contains(&account.name) {
547+
continue;
548+
}
549+
550+
let Ok((wallet, name)) = decompose_name(&account.name) else {
551+
warn!("Invalid account name {}", account.name);
552+
continue;
553+
};
554+
555+
match BlsPublicKey::try_from(account.public_key.as_slice()) {
556+
Ok(public_key) => {
557+
consensus_accounts.insert(
558+
public_key,
559+
Account::Simple(SimpleAccount {
560+
public_key,
561+
server: host.url.clone(),
562+
wallet,
563+
name,
564+
}),
565+
);
566+
}
567+
Err(_) => {
568+
warn!("Failed to parse public key for account {}", account.name);
569+
continue;
570+
}
571+
}
572+
}
573+
}
574+
575+
fn load_distributed_accounts(
576+
accounts: Vec<crate::proto::v1::DistributedAccount>,
577+
host: &DirkHostConfig,
578+
consensus_accounts: &mut HashMap<BlsPublicKey, Account>,
579+
) -> eyre::Result<()> {
580+
let host_name = host
581+
.server_name
582+
.clone()
583+
.or_else(|| host.url.host_str().map(String::from))
584+
.ok_or(eyre::eyre!("Host name not found for server {}", host.url))?;
585+
586+
for account in accounts {
587+
if !host.accounts.contains(&account.name) {
588+
continue;
589+
}
590+
591+
let Ok(public_key) = BlsPublicKey::try_from(account.composite_public_key.as_slice()) else {
592+
warn!("Failed to parse composite public key for account {}", account.name);
593+
continue;
594+
};
595+
596+
let Some(&Endpoint { id: participant_id, .. }) =
597+
account.participants.iter().find(|participant| participant.name == host_name)
598+
else {
599+
warn!("Host {host_name} not found as participant for account {}", account.name);
600+
continue;
601+
};
602+
603+
match consensus_accounts.get_mut(&public_key) {
604+
Some(Account::Distributed(DistributedAccount { participants, .. })) => {
605+
participants.insert(participant_id as u32, host.url.clone());
606+
}
607+
None => {
608+
let Ok((wallet, name)) = decompose_name(&account.name) else {
609+
warn!("Invalid account name {}", account.name);
610+
continue;
611+
};
612+
613+
let mut participants = HashMap::with_capacity(account.participants.len());
614+
participants.insert(participant_id as u32, host.url.clone());
615+
616+
consensus_accounts.insert(
617+
public_key,
618+
Account::Distributed(DistributedAccount {
619+
composite_public_key: public_key,
620+
participants,
621+
threshold: account.signing_threshold,
622+
wallet,
623+
name,
624+
}),
625+
);
626+
}
627+
Some(Account::Simple(_)) => {
628+
bail!("Distributed public key already exists for simple account");
629+
}
630+
}
631+
}
632+
633+
Ok(())
585634
}
586635

587-
pub fn aggregate_partial_signatures(
588-
partials: &[(BlsSignature, u32)],
589-
) -> eyre::Result<BlsSignature> {
636+
fn aggregate_partial_signatures(partials: &[(BlsSignature, u32)]) -> eyre::Result<BlsSignature> {
590637
// Deserialize partial signatures into G2 points
591638
let mut shares: HashMap<u32, G2Projective> = HashMap::new();
592639
for (sig, id) in partials {

0 commit comments

Comments
 (0)