From bfbf200652156fa9a1b8d081f0eb91522818e404 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jo=C3=A3o=20Lucas?= Date: Sat, 18 Apr 2026 15:57:52 -0300 Subject: [PATCH] perf: 11 allocation trims from hot-path audit MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Audit pass across the send, receive, retry, and signal paths — pattern matches the same shape as the previous PR's win (eager to_non_ad() / clone() then discarded, map without with_capacity, clone for consumer that only borrows). No behavior change in any of the spots. Per-message / per-receive: - message.rs::parse_message_info: stop cloning the entire `Device` just to read .pn and .lid. Read the two fields through the RwLock guard directly, mirroring the get_pn() getter pattern. - message.rs (batch skmsg decrypt loop): hoist sender_for_sk, sender_address, and sender_key_name out of `for payload in payloads` — all three are loop-invariant. Saves 3 allocations per extra payload in a batch (common for group status). Per-send: - send.rs (DM fanout): replace the `HashSet::insert(j.clone())` dedup-retain with `sort_dedup_by_device`. participant_list_hash sorts internally, so reordering is safe; now zero allocations for dedup. - wacore/src/send.rs (force-SKDM distribution): don't pre-allocate own_jid_to_check before checking presence in the list — compare by `&str` user, allocate only on the push branch. Saves 1 Jid per send on the "own already present" path (the common case). Per-retry: - retry.rs: drop the second processing_key.clone() — move it into the scopeguard closure instead of cloning twice. - retry.rs: replace `info.requester.user.clone()` with a direct borrow where the only consumer is `has_device(&str)`. - retry.rs: turn `recipient.clone().unwrap().to_non_ad()` into `recipient.as_ref()` + pattern match. Getters: - client.rs: get_push_name and get_lid were cloning the whole device_snapshot. Match get_pn: read the one field via persistence_manager.get_device_arc().read().await..clone(). - sender_keys.rs: borrow `own_lid_user` / `own_pn_user` as `&str` into the filter closure instead of cloning each `CompactString` up-front — the snapshot stays alive for the .collect() call. Other: - features/groups.rs::query_info: collapse two iterations of `group.participants` into a single move-based loop — saves one Jid clone per participant, plus one CompactString clone per LID entry with a PN mapping. - features/polls.rs::vote: replace `my_jid.to_non_ad() == poll_creator_jid.to_non_ad()` with `is_same_user_as` (the same fix shape from the reference PR). Cache `my_base = my_jid.to_non_ad()` so the voter_jid_str derivation and the equality check share one allocation. - features/polls.rs::aggregate_votes: inline the decryption helper to cache `creator_str` once instead of recomputing `poll_creator_jid.to_non_ad().to_string()` per voter. Added HashMap::with_capacity(votes.len()) while we're here. - wacore/src/prekeys.rs::parse_prekeys_response: HashMap::with_capacity(children.len()). - sender_key_device_cache.rs::from_db_rows: HashMap::with_capacity(rows.len()), HashSet ditto. 552 wacore lib tests + all per-crate suites remain green. --- src/client.rs | 18 ++++++++++++++---- src/client/sender_keys.rs | 22 +++++++++++++--------- src/features/groups.rs | 31 +++++++++++++++---------------- src/features/polls.rs | 32 +++++++++++++++++++++----------- src/message.rs | 25 +++++++++++++++---------- src/retry.rs | 19 +++++++++++-------- src/send.rs | 10 ++++------ src/sender_key_device_cache.rs | 4 ++-- wacore/src/prekeys.rs | 5 +++-- wacore/src/send.rs | 28 +++++++++++++++------------- 10 files changed, 113 insertions(+), 81 deletions(-) diff --git a/src/client.rs b/src/client.rs index 713821e9c..4db92934e 100644 --- a/src/client.rs +++ b/src/client.rs @@ -3581,8 +3581,13 @@ impl Client { } pub async fn get_push_name(&self) -> String { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; - device_snapshot.push_name.clone() + self.persistence_manager + .get_device_arc() + .await + .read() + .await + .push_name + .clone() } pub async fn get_pn(&self) -> Option { @@ -3596,8 +3601,13 @@ impl Client { } pub async fn get_lid(&self) -> Option { - let snapshot = self.persistence_manager.get_device_snapshot().await; - snapshot.lid.clone() + self.persistence_manager + .get_device_arc() + .await + .read() + .await + .lid + .clone() } pub(crate) async fn require_pn(&self) -> Result { diff --git a/src/client/sender_keys.rs b/src/client/sender_keys.rs index 2b24b869a..92ef97a16 100644 --- a/src/client/sender_keys.rs +++ b/src/client/sender_keys.rs @@ -15,22 +15,26 @@ impl Client { has_key: bool, exclude_own_devices: bool, ) -> Result<()> { - let (own_lid_user, own_pn_user) = if exclude_own_devices { - let snapshot = self.persistence_manager.get_device_snapshot().await; - ( - snapshot.lid.as_ref().map(|j| j.user.clone()), - snapshot.pn.as_ref().map(|j| j.user.clone()), - ) + let snapshot = if exclude_own_devices { + Some(self.persistence_manager.get_device_snapshot().await) } else { - (None, None) + None }; + let own_lid_user = snapshot + .as_ref() + .and_then(|s| s.lid.as_ref()) + .map(|j| j.user.as_str()); + let own_pn_user = snapshot + .as_ref() + .and_then(|s| s.pn.as_ref()) + .map(|j| j.user.as_str()); let device_ids: Vec = device_jids .iter() .filter(|jid| { !exclude_own_devices - || !(own_lid_user.as_deref().is_some_and(|u| u == jid.user) - || own_pn_user.as_deref().is_some_and(|u| u == jid.user)) + || !(own_lid_user.is_some_and(|u| u == jid.user) + || own_pn_user.is_some_and(|u| u == jid.user)) }) .map(ToString::to_string) .collect(); diff --git a/src/features/groups.rs b/src/features/groups.rs index d492c53f2..881f11157 100644 --- a/src/features/groups.rs +++ b/src/features/groups.rs @@ -193,22 +193,21 @@ impl<'a> Groups<'a> { let group = self.client.execute(GroupQueryIq::new(jid)).await?; - let participants: Vec = group.participants.iter().map(|p| p.jid.clone()).collect(); - - let lid_to_pn_map: HashMap = - if group.addressing_mode == AddressingMode::Lid { - group - .participants - .iter() - .filter_map(|p| { - p.phone_number - .as_ref() - .map(|pn| (p.jid.user.clone(), pn.clone())) - }) - .collect() - } else { - HashMap::new() - }; + // Single pass: move participants out and build lid_to_pn_map alongside. + let n = group.participants.len(); + let is_lid = group.addressing_mode == AddressingMode::Lid; + let mut participants: Vec = Vec::with_capacity(n); + let mut lid_to_pn_map: HashMap = if is_lid { + HashMap::with_capacity(n) + } else { + HashMap::new() + }; + for p in group.participants { + if is_lid && let Some(pn) = p.phone_number { + lid_to_pn_map.insert(p.jid.user.clone(), pn); + } + participants.push(p.jid); + } let mut info = GroupInfo::new(participants, group.addressing_mode); if !lid_to_pn_map.is_empty() { diff --git a/src/features/polls.rs b/src/features/polls.rs index b6c97f40c..16dbb51f4 100644 --- a/src/features/polls.rs +++ b/src/features/polls.rs @@ -118,7 +118,8 @@ impl<'a> Polls<'a> { .get_pn() .await .ok_or_else(|| anyhow!("Not logged in — cannot determine own JID"))?; - let voter_jid_str = my_jid.to_non_ad().to_string(); + let my_base = my_jid.to_non_ad(); + let voter_jid_str = my_base.to_string(); let creator_jid_str = poll_creator_jid.to_non_ad().to_string(); let selected_hashes: Vec> = option_names @@ -136,7 +137,7 @@ impl<'a> Polls<'a> { let (enc_payload, iv) = poll::encrypt_poll_vote(&selected_hashes, &key, poll_msg_id, &voter_jid_str)?; - let from_me = my_jid.to_non_ad() == poll_creator_jid.to_non_ad(); + let from_me = my_base.is_same_user_as(poll_creator_jid); let poll_update = wa::message::PollUpdateMessage { poll_creation_message_key: Some(wa::MessageKey { @@ -196,24 +197,33 @@ impl<'a> Polls<'a> { .map(|name| (poll::compute_option_hash(name), name.as_str())) .collect(); + // `creator_str` is invariant across voters; `decrypt_vote` used to + // recompute it per voter via `poll_creator_jid.to_non_ad().to_string()`. + let creator_str = poll_creator_jid.to_non_ad().to_string(); + // Last-vote-wins: each new vote from the same voter replaces the previous - let mut latest_votes: HashMap>> = HashMap::new(); + let mut latest_votes: HashMap>> = HashMap::with_capacity(votes.len()); for (voter_jid, enc_payload, enc_iv) in votes { - let voter_key = voter_jid.to_non_ad().to_string(); - match Self::decrypt_vote( - enc_payload, - enc_iv, + let voter_str = voter_jid.to_non_ad().to_string(); + let key = match poll::derive_vote_encryption_key( message_secret, poll_msg_id, - poll_creator_jid, - voter_jid, + &creator_str, + &voter_str, ) { + Ok(k) => k, + Err(e) => { + log::warn!("Failed to derive vote key for {voter_jid}: {e}"); + continue; + } + }; + match poll::decrypt_poll_vote(enc_payload, enc_iv, &key, poll_msg_id, &voter_str) { Ok(selected_hashes) => { if selected_hashes.is_empty() { // Empty selection = voter cleared their vote - latest_votes.remove(&voter_key); + latest_votes.remove(&voter_str); } else { - latest_votes.insert(voter_key, selected_hashes); + latest_votes.insert(voter_str, selected_hashes); } } Err(e) => { diff --git a/src/message.rs b/src/message.rs index 0cc6d9c1b..cd3e977c4 100644 --- a/src/message.rs +++ b/src/message.rs @@ -1098,17 +1098,18 @@ impl Client { } let mut adapter = self.signal_adapter().await; + // Always use bare sender for sender key operations. Real WA delivers + // skmsg with bare participant but pkmsg (SKDM) with device-qualified + // participant — normalizing to bare ensures consistent lookup. + // Hoisted out of the payload loop: all three are loop-invariant. + let sender_for_sk = info.source.sender.to_non_ad(); + let sender_address = sender_for_sk.to_protocol_address(); + let sender_key_name = make_sender_key_name(&info.source.chat, &sender_address); + for payload in payloads { let ciphertext = &payload.ciphertext[..]; let padding_version = payload.padding_version; - // Always use bare sender for sender key operations. Real WA delivers - // skmsg with bare participant but pkmsg (SKDM) with device-qualified - // participant — normalizing to bare ensures consistent lookup. - let sender_for_sk = info.source.sender.to_non_ad(); - let sender_address = sender_for_sk.to_protocol_address(); - let sender_key_name = make_sender_key_name(&info.source.chat, &sender_address); - log::debug!( "Looking up sender key for group {} with sender address {} (from sender JID: {})", info.source.chat, @@ -1434,10 +1435,14 @@ impl Client { &self, node: &wacore_binary::NodeRef<'_>, ) -> Result { - let device_snapshot = self.persistence_manager.get_device_snapshot().await; + let (own_pn, own_lid) = { + let arc = self.persistence_manager.get_device_arc().await; + let guard = arc.read().await; + (guard.pn.clone(), guard.lid.clone()) + }; let default_jid = Jid::default(); - let own_jid = device_snapshot.pn.as_ref().unwrap_or(&default_jid); - wacore::messages::parse_message_info(node, own_jid, device_snapshot.lid.as_ref()) + let own_jid = own_pn.as_ref().unwrap_or(&default_jid); + wacore::messages::parse_message_info(node, own_jid, own_lid.as_ref()) } pub(crate) async fn handle_app_state_sync_key_share( diff --git a/src/retry.rs b/src/retry.rs index 7f74692e5..53cdec9f5 100644 --- a/src/retry.rs +++ b/src/retry.rs @@ -133,10 +133,10 @@ fn resolve_retry_chat_info( let is_peer = own_pn.is_some_and(|pn| from.is_same_user_as(pn)) || own_lid.is_some_and(|lid| from.is_same_user_as(lid)); - let chat = if is_bot && recipient.is_some() { - recipient.clone().unwrap().to_non_ad() + let chat = if is_bot && let Some(r) = recipient.as_ref() { + r.to_non_ad() } else if is_peer { - match &recipient { + match recipient.as_ref() { Some(r) => r.to_non_ad(), // No recipient on peer retry — chat will be our own JID, // message lookup will likely fail. WA Web returns null here. @@ -233,13 +233,14 @@ impl Client { log::debug!("Ignoring retry for {processing_key}: a retry is already in progress."); return Ok(()); } + // processing_key isn't needed by name after this point — move it into + // the scopeguard instead of cloning again. let pending = Arc::clone(&self.pending_retries); - let guard_key = processing_key.clone(); let _guard = scopeguard::guard((), move |()| { pending .lock() .unwrap_or_else(|p| p.into_inner()) - .remove(&guard_key); + .remove(&processing_key); }); let (original_msg, alt_chat) = match self.take_recent_message(&info.chat, &message_id).await @@ -282,11 +283,13 @@ impl Client { }; let sender_device_id = info.requester.device() as u32; - let sender_user = info.requester.user.clone(); - if !self.has_device(&sender_user, sender_device_id).await { + if !self + .has_device(&info.requester.user, sender_device_id) + .await + { warn!( "handle_retry_receipt: device not found for device={}, user={}", - sender_device_id, sender_user + sender_device_id, info.requester.user ); return Ok(()); } diff --git a/src/send.rs b/src/send.rs index 50b9dab26..376643e11 100644 --- a/src/send.rs +++ b/src/send.rs @@ -1154,12 +1154,10 @@ impl Client { !is_sender }); - // Dedup for self-DMs: recipient and own device lists overlap - // when sending to own account (WA Web uses Map keyed by toString) - { - let mut seen = std::collections::HashSet::with_capacity(all_dm_jids.len()); - all_dm_jids.retain(|j| seen.insert(j.clone())); - } + // Dedup for self-DMs: recipient and own device lists overlap when + // sending to own account. `participant_list_hash` sorts internally, + // so reordering here is safe. + wacore::types::jid::sort_dedup_by_device(&mut all_dm_jids); self.ensure_e2e_sessions(&all_dm_jids).await?; diff --git a/src/sender_key_device_cache.rs b/src/sender_key_device_cache.rs index 516e88874..4131e44ff 100644 --- a/src/sender_key_device_cache.rs +++ b/src/sender_key_device_cache.rs @@ -19,8 +19,8 @@ pub(crate) struct SenderKeyDeviceMap { impl SenderKeyDeviceMap { pub fn from_db_rows(rows: &[(String, bool)]) -> Self { - let mut devices: HashMap, HashMap> = HashMap::new(); - let mut forgotten_users = HashSet::new(); + let mut devices: HashMap, HashMap> = HashMap::with_capacity(rows.len()); + let mut forgotten_users = HashSet::with_capacity(rows.len() / 4); for (jid_str, has_key) in rows { match jid_str.parse::() { diff --git a/wacore/src/prekeys.rs b/wacore/src/prekeys.rs index d03132717..14074860d 100644 --- a/wacore/src/prekeys.rs +++ b/wacore/src/prekeys.rs @@ -165,8 +165,9 @@ impl PreKeyUtils { .get_optional_child("list") .ok_or_else(|| anyhow::anyhow!(" not found in pre-key response"))?; - let mut bundles = HashMap::new(); - for user_node_ref in list_node.children().unwrap_or_default() { + let children = list_node.children().unwrap_or_default(); + let mut bundles = HashMap::with_capacity(children.len()); + for user_node_ref in children { if user_node_ref.tag != "user" { continue; } diff --git a/wacore/src/send.rs b/wacore/src/send.rs index 9750c0861..463a9130f 100644 --- a/wacore/src/send.rs +++ b/wacore/src/send.rs @@ -1090,21 +1090,23 @@ pub async fn prepare_group_stanza< }) .collect(); - // Determine what JID to check for - use phone number if we're in LID mode and have a mapping - let own_jid_to_check = if own_base_jid.is_lid() { - group_info - .phone_jid_for_lid_user(&own_base_jid.user) - .map(|pn| pn.to_non_ad()) - .unwrap_or_else(|| own_base_jid.clone()) + // Determine what user to check for — use the PN user when own is LID + // and we have a mapping. Keeping this as a borrow avoids allocating a + // throwaway Jid when own is already in the list. + let own_pn_mapping = if own_base_jid.is_lid() { + group_info.phone_jid_for_lid_user(&own_base_jid.user) } else { - own_base_jid.clone() + None }; - - if !jids_to_resolve - .iter() - .any(|participant| participant.is_same_user_as(&own_jid_to_check)) - { - jids_to_resolve.push(own_jid_to_check); + let own_check_user = own_pn_mapping + .map(|pn| pn.user.as_str()) + .unwrap_or(own_base_jid.user.as_str()); + + if !jids_to_resolve.iter().any(|p| p.user == own_check_user) { + jids_to_resolve.push(match own_pn_mapping { + Some(pn) => pn.to_non_ad(), + None => own_base_jid.clone(), + }); } crate::types::jid::sort_dedup_by_user(&mut jids_to_resolve);