perf!: reduce startup allocations by ~21% and peak heap by ~31% - #525
Merged
Conversation
📝 WalkthroughWalkthroughBatch prekey load/store APIs and implementations were added; prekey validation now uses batch loading. Request builders and IQ construction were refactored to take owned values/iterators to avoid clones. HTTP client agent builder now uses smaller buffers and reduced connection pooling across all code paths. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as Client
participant SignalStore as SignalStore (trait impl)
participant DB as Storage (Sqlite / InMemory)
rect rgba(100,150,240,0.5)
Client->>SignalStore: validate_digest_key(requested_ids)
end
rect rgba(120,200,80,0.5)
SignalStore->>DB: load_prekeys_batch(ids)
DB-->>SignalStore: Vec<(id, bytes)>
SignalStore-->>Client: Vec<(id, bytes)> or empty/missing info
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
jlucaso1
force-pushed
the
perf/reduce-startup-allocations
branch
2 times, most recently
from
April 13, 2026 05:31
f027018 to
5b7d82a
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bd0e610118
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+1084
to
+1086
| Some(mut devices) => { | ||
| devices.retain(|j| !j.is_hosted()); | ||
| devices |
There was a problem hiding this comment.
If recipient_cached exists but is empty after hosted-device filtering, this branch returns an empty recipient target set and never adds recipient_bare. In that case prepare_dm_stanza can build a stanza that encrypts only for own devices (or no devices), so the remote user may not receive the message even though send returns success. This is a regression from the previous behavior that always included bare recipient fanout, and it can happen with stale/partial registry records (including hosted-only entries).
Useful? React with 👍 / 👎.
Benchmark Results59 unchanged benchmark(s)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@src/prekeys.rs`:
- Around line 323-331: The check currently compares loaded_map.len() (deduped)
to response.prekey_ids.len() (may contain duplicates) causing false misses;
change the comparison to dedupe response.prekey_ids first (e.g., build a HashSet
from response.prekey_ids) and compare loaded_map.len() against that deduped set
size, while continuing to use the original response.prekey_ids list for digest
ordering and iteration; update the branch that logs and returns to use the
deduped count so repeated IDs don't trigger a false missing-local-prekeys
warning in the code paths around loaded_map and response.prekey_ids.
In `@src/send.rs`:
- Around line 1068-1104: The code discards the successful results from
get_user_devices() and always re-reads the registry, which can cause fallback to
bare-JID fanout or omit own devices; change the logic to use the Vec<Jid>
returned by get_user_devices() when it succeeds instead of immediately calling
get_devices_from_registry() again: when calling get_user_devices(&[to]) or
get_user_devices(own_jid) capture and use its returned Vec<Jid> to populate
recipient_cached and own_cached (or directly append into all_dm_jids), preserve
the existing hosted-device filtering (j.is_hosted()), exclude the sender using
the existing own_lid check, and deduplicate the final all_dm_jids before calling
ensure_e2e_sessions() so you don’t rely solely on registry visibility.
In `@wacore/src/iq/props.rs`:
- Around line 198-207: The parsing currently prefers Self::Experiment when both
discriminators are present; update the conditional in the prop discriminator
logic to explicitly detect the ambiguous case where optional_attr(node,
"config_code").is_some() && optional_attr(node, "event_code").is_some() and
return an Err (include node.attrs for context) instead of falling through to
AbProp::try_from_node_ref; otherwise proceed to construct Self::Experiment via
AbProp::try_from_node_ref or Self::Sampling via SamplingProp::try_from_node_ref
as before.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: bfbc90c0-fb68-493f-b5cc-2f58ff32e6a5
📒 Files selected for processing (12)
http_clients/ureq-client/src/lib.rssrc/prekeys.rssrc/request.rssrc/send.rsstorages/sqlite-storage/src/sqlite_store.rswacore/src/iq/prekeys.rswacore/src/iq/props.rswacore/src/prekeys.rswacore/src/request.rswacore/src/send.rswacore/src/store/in_memory.rswacore/src/store/traits.rs
Comment on lines
+1068
to
+1104
| // Local registry first; network warm only on miss to avoid | ||
| // unnecessary LID-migration side effects from get_user_devices | ||
| let mut recipient_cached = self.get_devices_from_registry(&recipient_bare).await; | ||
| if recipient_cached.is_none() { | ||
| let _ = self.get_user_devices(std::slice::from_ref(&to)).await; | ||
| recipient_cached = self.get_devices_from_registry(&recipient_bare).await; | ||
| } | ||
|
|
||
| let mut own_cached = self.get_devices_from_registry(own_jid).await; | ||
| if own_cached.is_none() { | ||
| let _ = self.get_user_devices(std::slice::from_ref(own_jid)).await; | ||
| own_cached = self.get_devices_from_registry(own_jid).await; | ||
| } | ||
|
|
||
| // Build device list, filter hosted in-place, reuse Vecs | ||
| let mut all_dm_jids = match recipient_cached { | ||
| Some(mut devices) => { | ||
| devices.retain(|j| !j.is_hosted()); | ||
| devices | ||
| } | ||
| // No record at all — bare JID, server handles fanout | ||
| None => vec![recipient_bare], | ||
| }; | ||
|
|
||
| if let Some(mut own_devices) = own_cached { | ||
| own_devices.retain(|j| !j.is_hosted()); | ||
| all_dm_jids.append(&mut own_devices); | ||
| } | ||
|
|
||
| let mut all_dm_jids = Vec::with_capacity(1 + own_devices.len()); | ||
| all_dm_jids.push(recipient_bare); | ||
| all_dm_jids.extend(own_devices); | ||
| // Exclude exact sender device (WA Web: isMeDevice in getFanOutList) | ||
| // so ensure_e2e_sessions never creates a self-session | ||
| let own_lid = device_snapshot.lid.as_ref(); | ||
| all_dm_jids.retain(|j| { | ||
| let is_sender = (j.is_same_user_as(own_jid) && j.device == own_jid.device) | ||
| || own_lid.is_some_and(|lid| j.is_same_user_as(lid) && j.device == lid.device); | ||
| !is_sender | ||
| }); |
There was a problem hiding this comment.
Don't discard successful get_user_devices() results here.
get_user_devices() already returns the resolved device list, but this path ignores that Vec<Jid> and immediately re-reads the registry. If the refresh succeeds but the cache/DB write is not visible yet, recipients fall back to bare-JID fanout and own companions are silently omitted from DSM fanout. Build all_dm_jids from the returned lists directly, then dedup before ensure_e2e_sessions().
🛠️ Suggested fix
- let mut recipient_cached = self.get_devices_from_registry(&recipient_bare).await;
- if recipient_cached.is_none() {
- let _ = self.get_user_devices(std::slice::from_ref(&to)).await;
- recipient_cached = self.get_devices_from_registry(&recipient_bare).await;
- }
-
- let mut own_cached = self.get_devices_from_registry(own_jid).await;
- if own_cached.is_none() {
- let _ = self.get_user_devices(std::slice::from_ref(own_jid)).await;
- own_cached = self.get_devices_from_registry(own_jid).await;
- }
+ let mut recipient_cached = self.get_devices_from_registry(&recipient_bare).await;
+ if recipient_cached.is_none() {
+ recipient_cached = self
+ .get_user_devices(std::slice::from_ref(&to))
+ .await
+ .ok();
+ }
+
+ let mut own_cached = self.get_devices_from_registry(own_jid).await;
+ if own_cached.is_none() {
+ own_cached = self
+ .get_user_devices(std::slice::from_ref(own_jid))
+ .await
+ .ok();
+ }
@@
- if let Some(mut own_devices) = own_cached {
+ if let Some(mut own_devices) = own_cached {
own_devices.retain(|j| !j.is_hosted());
all_dm_jids.append(&mut own_devices);
}
+
+ wacore::types::jid::sort_dedup_by_device(&mut all_dm_jids);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| // Local registry first; network warm only on miss to avoid | |
| // unnecessary LID-migration side effects from get_user_devices | |
| let mut recipient_cached = self.get_devices_from_registry(&recipient_bare).await; | |
| if recipient_cached.is_none() { | |
| let _ = self.get_user_devices(std::slice::from_ref(&to)).await; | |
| recipient_cached = self.get_devices_from_registry(&recipient_bare).await; | |
| } | |
| let mut own_cached = self.get_devices_from_registry(own_jid).await; | |
| if own_cached.is_none() { | |
| let _ = self.get_user_devices(std::slice::from_ref(own_jid)).await; | |
| own_cached = self.get_devices_from_registry(own_jid).await; | |
| } | |
| // Build device list, filter hosted in-place, reuse Vecs | |
| let mut all_dm_jids = match recipient_cached { | |
| Some(mut devices) => { | |
| devices.retain(|j| !j.is_hosted()); | |
| devices | |
| } | |
| // No record at all — bare JID, server handles fanout | |
| None => vec![recipient_bare], | |
| }; | |
| if let Some(mut own_devices) = own_cached { | |
| own_devices.retain(|j| !j.is_hosted()); | |
| all_dm_jids.append(&mut own_devices); | |
| } | |
| let mut all_dm_jids = Vec::with_capacity(1 + own_devices.len()); | |
| all_dm_jids.push(recipient_bare); | |
| all_dm_jids.extend(own_devices); | |
| // Exclude exact sender device (WA Web: isMeDevice in getFanOutList) | |
| // so ensure_e2e_sessions never creates a self-session | |
| let own_lid = device_snapshot.lid.as_ref(); | |
| all_dm_jids.retain(|j| { | |
| let is_sender = (j.is_same_user_as(own_jid) && j.device == own_jid.device) | |
| || own_lid.is_some_and(|lid| j.is_same_user_as(lid) && j.device == lid.device); | |
| !is_sender | |
| }); | |
| // Local registry first; network warm only on miss to avoid | |
| // unnecessary LID-migration side effects from get_user_devices | |
| let mut recipient_cached = self.get_devices_from_registry(&recipient_bare).await; | |
| if recipient_cached.is_none() { | |
| recipient_cached = self | |
| .get_user_devices(std::slice::from_ref(&to)) | |
| .await | |
| .ok(); | |
| } | |
| let mut own_cached = self.get_devices_from_registry(own_jid).await; | |
| if own_cached.is_none() { | |
| own_cached = self | |
| .get_user_devices(std::slice::from_ref(own_jid)) | |
| .await | |
| .ok(); | |
| } | |
| // Build device list, filter hosted in-place, reuse Vecs | |
| let mut all_dm_jids = match recipient_cached { | |
| Some(mut devices) => { | |
| devices.retain(|j| !j.is_hosted()); | |
| devices | |
| } | |
| // No record at all — bare JID, server handles fanout | |
| None => vec![recipient_bare], | |
| }; | |
| if let Some(mut own_devices) = own_cached { | |
| own_devices.retain(|j| !j.is_hosted()); | |
| all_dm_jids.append(&mut own_devices); | |
| } | |
| wacore::types::jid::sort_dedup_by_device(&mut all_dm_jids); | |
| // Exclude exact sender device (WA Web: isMeDevice in getFanOutList) | |
| // so ensure_e2e_sessions never creates a self-session | |
| let own_lid = device_snapshot.lid.as_ref(); | |
| all_dm_jids.retain(|j| { | |
| let is_sender = (j.is_same_user_as(own_jid) && j.device == own_jid.device) | |
| || own_lid.is_some_and(|lid| j.is_same_user_as(lid) && j.device == lid.device); | |
| !is_sender | |
| }); |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/send.rs` around lines 1068 - 1104, The code discards the successful
results from get_user_devices() and always re-reads the registry, which can
cause fallback to bare-JID fanout or omit own devices; change the logic to use
the Vec<Jid> returned by get_user_devices() when it succeeds instead of
immediately calling get_devices_from_registry() again: when calling
get_user_devices(&[to]) or get_user_devices(own_jid) capture and use its
returned Vec<Jid> to populate recipient_cached and own_cached (or directly
append into all_dm_jids), preserve the existing hosted-device filtering
(j.is_hosted()), exclude the sender using the existing own_lid check, and
deduplicate the final all_dm_jids before calling ensure_e2e_sessions() so you
don’t rely solely on registry visibility.
Heaptrack-guided optimizations targeting connection startup hot paths: - build_iq_node: take InfoQuery by value, eliminating deep Node tree clones (was 17K alloc calls for prekey upload alone) - build_upload_prekeys_request: accept iterator instead of slice, removing intermediate Vec collect and per-key public_bytes clone - AbPropConfig: check discriminating attribute before parsing to avoid double-parse allocations on every prop node - AbProp config_value: use into_owned() instead of to_string() on Cow - SignalStore: add has_session() with default impl; InMemoryBackend and SqliteStore override with contains_key/SELECT EXISTS to avoid loading full session blobs - SignalStore: add load_prekeys_batch() and store_prekeys_batch() overrides for InMemoryBackend (single lock) and SqliteStore (WHERE IN query), eliminating 2,400+ per-key async_trait box allocs - validate_digest_key: batch-load prekeys instead of 812 individual load_prekey calls - ureq: configure 16KB buffers (down from 128KB default) and reduce idle connection pool; media uses streaming I/O so throughput is unaffected Benchmark results (1-message mock server, heaptrack): Before: 29,869 allocs, 3.17 MB peak heap, 22.02 MB RSS After: 23,622 allocs, 2.19 MB peak heap, 20.06 MB RSS BREAKING: build_iq_node takes InfoQuery by value; build_upload_prekeys_request takes impl IntoIterator.
jlucaso1
force-pushed
the
perf/reduce-startup-allocations
branch
from
April 13, 2026 05:41
5b7d82a to
a298676
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
wacore/src/iq/props.rs (1)
198-208:⚠️ Potential issue | 🟡 MinorHandle ambiguous discriminators explicitly (
config_code+event_code).A malformed
<prop>containing bothconfig_codeandevent_codeattributes is silently treated asExperiment(line 199-200). This should be rejected explicitly to avoid misclassification and ensure defensive parsing.🛡️ Proposed fix to detect and reject ambiguous cases
- // Check discriminating attribute to avoid double-parse allocations - if optional_attr(node, "config_code").is_some() { - Ok(Self::Experiment(AbProp::try_from_node_ref(node)?)) - } else if optional_attr(node, "event_code").is_some() { - Ok(Self::Sampling(SamplingProp::try_from_node_ref(node)?)) - } else { - Err(anyhow::anyhow!( - "prop has neither config_code nor event_code (attrs: {:?})", - node.attrs - )) - } + // Check discriminating attribute to avoid double-parse allocations + let has_config_code = optional_attr(node, "config_code").is_some(); + let has_event_code = optional_attr(node, "event_code").is_some(); + + match (has_config_code, has_event_code) { + (true, false) => Ok(Self::Experiment(AbProp::try_from_node_ref(node)?)), + (false, true) => Ok(Self::Sampling(SamplingProp::try_from_node_ref(node)?)), + (true, true) => Err(anyhow::anyhow!( + "prop has both config_code and event_code (attrs: {:?})", + node.attrs + )), + (false, false) => Err(anyhow::anyhow!( + "prop has neither config_code nor event_code (attrs: {:?})", + node.attrs + )), + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/src/iq/props.rs` around lines 198 - 208, The current discriminator logic silently prefers Experiment when both attributes exist; update the branch in the constructor where optional_attr(node, "config_code") and optional_attr(node, "event_code") are checked to explicitly detect the ambiguous case (both present) and return an Err with a clear message including node.attrs, rather than falling through to Experiment. Modify the flow around Self::Experiment(AbProp::try_from_node_ref(node)?) and Self::Sampling(SamplingProp::try_from_node_ref(node)?) so you first check for both attributes and error, then check singletons to call try_from_node_ref for AbProp or SamplingProp accordingly.src/prekeys.rs (1)
323-331:⚠️ Potential issue | 🟡 MinorDeduplicate the server IDs before the missing-prekey check.
Line 326 compares
loaded_map.len()againstresponse.prekey_ids.len(), butloaded_mapis already deduped by key. If the digest list repeats an ID, this branch will still log a false miss and skip validation even though the local prekey exists.[suggested fix]
Proposed patch
// Build a lookup so we preserve the server-requested order let loaded_map: std::collections::HashMap<u32, Vec<u8>> = loaded.into_iter().collect(); + let unique_prekey_ids: std::collections::HashSet<u32> = + response.prekey_ids.iter().copied().collect(); - if loaded_map.len() < response.prekey_ids.len() { + if loaded_map.len() < unique_prekey_ids.len() { log::warn!( "digestKey: missing {} local prekeys, skipping", - response.prekey_ids.len() - loaded_map.len() + unique_prekey_ids.len() - loaded_map.len() ); return Ok(()); }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/prekeys.rs` around lines 323 - 331, The missing-prekey check compares loaded_map.len() to response.prekey_ids.len(), but response.prekey_ids may contain duplicates; deduplicate the server IDs before that comparison by creating a unique set (e.g., a HashSet or a deduped Vec) from response.prekey_ids and use unique_ids.len() in the comparison and log message instead of response.prekey_ids.len(); update references around loaded_map and the warning log so you count unique server-requested IDs when deciding to early-return from digestKey.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@storages/sqlite-storage/src/sqlite_store.rs`:
- Around line 1369-1389: The batch loader currently converts the input ids and
returns DB rows in database order, losing caller order and duplicates; modify
load_prekeys_batch so after loading rows (variable rows) you build a lookup
(e.g., HashMap<i32, Vec<u8>> or HashMap<i32, Vec<u8>>) from the returned rows
and then iterate the original input order (preserve the original ids slice
before converting to i32), pushing an entry for each requested id in request
order and repeating entries for duplicate requested ids when present in the DB;
ensure you move the preserved original ids into the semaphore closure and use
that to assemble the Vec<(u32, Vec<u8>)> so behavior matches the default
SignalStore::load_prekeys_batch ordering and duplicate semantics.
---
Duplicate comments:
In `@src/prekeys.rs`:
- Around line 323-331: The missing-prekey check compares loaded_map.len() to
response.prekey_ids.len(), but response.prekey_ids may contain duplicates;
deduplicate the server IDs before that comparison by creating a unique set
(e.g., a HashSet or a deduped Vec) from response.prekey_ids and use
unique_ids.len() in the comparison and log message instead of
response.prekey_ids.len(); update references around loaded_map and the warning
log so you count unique server-requested IDs when deciding to early-return from
digestKey.
In `@wacore/src/iq/props.rs`:
- Around line 198-208: The current discriminator logic silently prefers
Experiment when both attributes exist; update the branch in the constructor
where optional_attr(node, "config_code") and optional_attr(node, "event_code")
are checked to explicitly detect the ambiguous case (both present) and return an
Err with a clear message including node.attrs, rather than falling through to
Experiment. Modify the flow around
Self::Experiment(AbProp::try_from_node_ref(node)?) and
Self::Sampling(SamplingProp::try_from_node_ref(node)?) so you first check for
both attributes and error, then check singletons to call try_from_node_ref for
AbProp or SamplingProp accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 5de42919-bc0c-43b4-8138-99a646644939
📒 Files selected for processing (10)
http_clients/ureq-client/src/lib.rssrc/prekeys.rssrc/request.rsstorages/sqlite-storage/src/sqlite_store.rswacore/src/iq/prekeys.rswacore/src/iq/props.rswacore/src/prekeys.rswacore/src/request.rswacore/src/store/in_memory.rswacore/src/store/traits.rs
Comment on lines
+1369
to
+1389
| async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> { | ||
| if ids.is_empty() { | ||
| return Ok(Vec::new()); | ||
| } | ||
| let pool = self.pool.clone(); | ||
| let device_id = self.device_id; | ||
| let ids: Vec<i32> = ids.iter().map(|&id| id as i32).collect(); | ||
| self.with_semaphore(move || -> Result<Vec<(u32, Vec<u8>)>> { | ||
| let mut conn = pool | ||
| .get() | ||
| .map_err(|e| StoreError::Connection(e.to_string()))?; | ||
| let rows: Vec<(i32, Vec<u8>)> = prekeys::table | ||
| .select((prekeys::id, prekeys::key)) | ||
| .filter(prekeys::id.eq_any(&ids)) | ||
| .filter(prekeys::device_id.eq(device_id)) | ||
| .load(&mut conn) | ||
| .map_err(|e| StoreError::Database(e.to_string()))?; | ||
| Ok(rows.into_iter().map(|(id, key)| (id as u32, key)).collect()) | ||
| }) | ||
| .await | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Preserve caller order in the SQLite batch path.
eq_any(&ids) returns matching rows in database order and only once per stored ID, while the default SignalStore::load_prekeys_batch implementation iterates the input slice and can return repeated IDs in request order. That makes this new API behave differently by backend.
Proposed patch
async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> {
if ids.is_empty() {
return Ok(Vec::new());
}
let pool = self.pool.clone();
let device_id = self.device_id;
- let ids: Vec<i32> = ids.iter().map(|&id| id as i32).collect();
+ let requested_ids: Vec<u32> = ids.to_vec();
+ let ids: Vec<i32> = requested_ids.iter().map(|&id| id as i32).collect();
self.with_semaphore(move || -> Result<Vec<(u32, Vec<u8>)>> {
let mut conn = pool
.get()
.map_err(|e| StoreError::Connection(e.to_string()))?;
let rows: Vec<(i32, Vec<u8>)> = prekeys::table
@@
.filter(prekeys::device_id.eq(device_id))
.load(&mut conn)
.map_err(|e| StoreError::Database(e.to_string()))?;
- Ok(rows.into_iter().map(|(id, key)| (id as u32, key)).collect())
+ let row_map: std::collections::HashMap<u32, Vec<u8>> =
+ rows.into_iter().map(|(id, key)| (id as u32, key)).collect();
+ Ok(requested_ids
+ .into_iter()
+ .filter_map(|id| row_map.get(&id).cloned().map(|key| (id, key)))
+ .collect())
})
.await
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> { | |
| if ids.is_empty() { | |
| return Ok(Vec::new()); | |
| } | |
| let pool = self.pool.clone(); | |
| let device_id = self.device_id; | |
| let ids: Vec<i32> = ids.iter().map(|&id| id as i32).collect(); | |
| self.with_semaphore(move || -> Result<Vec<(u32, Vec<u8>)>> { | |
| let mut conn = pool | |
| .get() | |
| .map_err(|e| StoreError::Connection(e.to_string()))?; | |
| let rows: Vec<(i32, Vec<u8>)> = prekeys::table | |
| .select((prekeys::id, prekeys::key)) | |
| .filter(prekeys::id.eq_any(&ids)) | |
| .filter(prekeys::device_id.eq(device_id)) | |
| .load(&mut conn) | |
| .map_err(|e| StoreError::Database(e.to_string()))?; | |
| Ok(rows.into_iter().map(|(id, key)| (id as u32, key)).collect()) | |
| }) | |
| .await | |
| } | |
| async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> { | |
| if ids.is_empty() { | |
| return Ok(Vec::new()); | |
| } | |
| let pool = self.pool.clone(); | |
| let device_id = self.device_id; | |
| let requested_ids: Vec<u32> = ids.to_vec(); | |
| let ids: Vec<i32> = requested_ids.iter().map(|&id| id as i32).collect(); | |
| self.with_semaphore(move || -> Result<Vec<(u32, Vec<u8>)>> { | |
| let mut conn = pool | |
| .get() | |
| .map_err(|e| StoreError::Connection(e.to_string()))?; | |
| let rows: Vec<(i32, Vec<u8>)> = prekeys::table | |
| .select((prekeys::id, prekeys::key)) | |
| .filter(prekeys::id.eq_any(&ids)) | |
| .filter(prekeys::device_id.eq(device_id)) | |
| .load(&mut conn) | |
| .map_err(|e| StoreError::Database(e.to_string()))?; | |
| let row_map: std::collections::HashMap<u32, Vec<u8>> = | |
| rows.into_iter().map(|(id, key)| (id as u32, key)).collect(); | |
| Ok(requested_ids | |
| .into_iter() | |
| .filter_map(|id| row_map.get(&id).cloned().map(|key| (id, key))) | |
| .collect()) | |
| }) | |
| .await | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@storages/sqlite-storage/src/sqlite_store.rs` around lines 1369 - 1389, The
batch loader currently converts the input ids and returns DB rows in database
order, losing caller order and duplicates; modify load_prekeys_batch so after
loading rows (variable rows) you build a lookup (e.g., HashMap<i32, Vec<u8>> or
HashMap<i32, Vec<u8>>) from the returned rows and then iterate the original
input order (preserve the original ids slice before converting to i32), pushing
an entry for each requested id in request order and repeating entries for
duplicate requested ids when present in the DB; ensure you move the preserved
original ids into the semaphore closure and use that to assemble the Vec<(u32,
Vec<u8>)> so behavior matches the default SignalStore::load_prekeys_batch
ordering and duplicate semantics.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/request.rs (1)
138-152: 🧹 Nitpick | 🔵 TrivialAvoid cloning
query.idnow thatsend_iqowns the query.This still allocates a fresh
Stringfor caller-supplied IDs before cloning it again for waiter registration / node building. Since the query is already owned here, take the ID out ofqueryinstead.♻️ Suggested change
- pub async fn send_iq( - &self, - query: InfoQuery<'_>, - ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> { + pub async fn send_iq( + &self, + mut query: InfoQuery<'_>, + ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> { @@ - let req_id = query - .id - .clone() - .unwrap_or_else(|| self.generate_request_id()); + let req_id = query.id.take().unwrap_or_else(|| self.generate_request_id());🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/request.rs` around lines 138 - 152, The code currently clones query.id unnecessarily even though send_iq owns query; replace query.id.clone().unwrap_or_else(...) with taking the ID out of the owned query (e.g. make query mutable if needed and use query.id.take().unwrap_or_else(|| self.generate_request_id())) so you consume the Option<String> instead of allocating a new String, then use that req_id for registering the waiter and for node building (keep any remaining .clone() only where a second owned copy is required).
♻️ Duplicate comments (1)
storages/sqlite-storage/src/sqlite_store.rs (1)
1369-1389:⚠️ Potential issue | 🟠 MajorPreserve caller order in the SQLite batch loader.
Line 1386 returns rows in database order and collapses repeated requested IDs, while the default
SignalStore::load_prekeys_batch()implementation walks the input slice and repeats entries for duplicate IDs. That makes this API backend-dependent. Rebuild the result from the original request list after querying.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@storages/sqlite-storage/src/sqlite_store.rs` around lines 1369 - 1389, The batch loader currently returns rows in DB order and collapses duplicates; to preserve caller order and duplicates, capture the original ids slice into a Vec<u32> (e.g., original_ids) before converting to Vec<i32>, then after loading rows inside load_prekeys_batch build a HashMap from fetched rows (id -> key) and iterate original_ids, pushing (id, key.clone()) into the result for each original id that exists in the map so duplicates are preserved; ensure you move original_ids into the closure used by with_semaphore and use key.clone() when repeating keys.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@wacore/src/iq/props.rs`:
- Around line 192-215: Add regression tests that exercise the two error branches
in the prop parser: one where both "config_code" and "event_code" attributes are
present and one where neither is present. Construct a node with tag "prop" and
attrs containing both keys and assert that Prop parsing returns an error
mentioning both attributes (exercise the has_config && has_event branch which
currently returns Err with node.attrs); likewise construct a node with tag
"prop" and no discriminator attrs and assert parsing returns the "neither" error
(exercise the else branch). Use the same entry point exercised by the code (the
Prop parsing / TryFrom node entry such as try_from_node_ref or the public
Prop->try_from_node API) so the tests fail if those branches regress.
In `@wacore/src/store/traits.rs`:
- Around line 129-139: Update the doc comment for the default implementation of
load_prekeys_batch to explicitly state its contract: it preserves the input
order and retains duplicates from the ids slice, includes one entry for each id
only if that id exists in storage (missing ids are omitted), and returns a
Vec<(u32, Vec<u8>)> in the same order as the input with duplicates repeated when
present; mention that custom backend implementations should match this
ordering/duplicate behavior to remain compatible with callers relying on it.
---
Outside diff comments:
In `@src/request.rs`:
- Around line 138-152: The code currently clones query.id unnecessarily even
though send_iq owns query; replace query.id.clone().unwrap_or_else(...) with
taking the ID out of the owned query (e.g. make query mutable if needed and use
query.id.take().unwrap_or_else(|| self.generate_request_id())) so you consume
the Option<String> instead of allocating a new String, then use that req_id for
registering the waiter and for node building (keep any remaining .clone() only
where a second owned copy is required).
---
Duplicate comments:
In `@storages/sqlite-storage/src/sqlite_store.rs`:
- Around line 1369-1389: The batch loader currently returns rows in DB order and
collapses duplicates; to preserve caller order and duplicates, capture the
original ids slice into a Vec<u32> (e.g., original_ids) before converting to
Vec<i32>, then after loading rows inside load_prekeys_batch build a HashMap from
fetched rows (id -> key) and iterate original_ids, pushing (id, key.clone())
into the result for each original id that exists in the map so duplicates are
preserved; ensure you move original_ids into the closure used by with_semaphore
and use key.clone() when repeating keys.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 00158dcb-1793-49e2-b27e-2e6703887d53
📒 Files selected for processing (10)
http_clients/ureq-client/src/lib.rssrc/prekeys.rssrc/request.rsstorages/sqlite-storage/src/sqlite_store.rswacore/src/iq/prekeys.rswacore/src/iq/props.rswacore/src/prekeys.rswacore/src/request.rswacore/src/store/in_memory.rswacore/src/store/traits.rs
Comment on lines
+192
to
+215
| use crate::iq::node::optional_attr; | ||
|
|
||
| if node.tag != "prop" { | ||
| return Err(anyhow::anyhow!("expected <prop>, got <{}>", node.tag)); | ||
| } | ||
|
|
||
| let experiment = AbProp::try_from_node_ref(node); | ||
| if let Ok(prop) = experiment { | ||
| return Ok(Self::Experiment(prop)); | ||
| } | ||
|
|
||
| let sampling = SamplingProp::try_from_node_ref(node); | ||
| if let Ok(prop) = sampling { | ||
| return Ok(Self::Sampling(prop)); | ||
| // Check discriminating attribute to avoid double-parse allocations | ||
| let has_config = optional_attr(node, "config_code").is_some(); | ||
| let has_event = optional_attr(node, "event_code").is_some(); | ||
|
|
||
| if has_config && has_event { | ||
| Err(anyhow::anyhow!( | ||
| "prop has both config_code and event_code (attrs: {:?})", | ||
| node.attrs | ||
| )) | ||
| } else if has_config { | ||
| Ok(Self::Experiment(AbProp::try_from_node_ref(node)?)) | ||
| } else if has_event { | ||
| Ok(Self::Sampling(SamplingProp::try_from_node_ref(node)?)) | ||
| } else { | ||
| Err(anyhow::anyhow!( | ||
| "prop has neither config_code nor event_code (attrs: {:?})", | ||
| node.attrs | ||
| )) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Add regression tests for the new discriminator error branches.
The parser logic is improved, but there’s no direct test coverage for the (config_code + event_code) and (neither) error paths.
Proposed test additions
@@
fn test_sampling_prop_protocol_node_round_trip() {
let prop = SamplingProp {
event_code: 5138,
sampling_weight: -1,
};
let node = prop.clone().into_node();
let parsed = SamplingProp::try_from_node(&node).unwrap();
assert_eq!(parsed.event_code, prop.event_code);
assert_eq!(parsed.sampling_weight, prop.sampling_weight);
}
+
+ #[test]
+ fn test_ab_prop_config_rejects_ambiguous_discriminator() {
+ let node = NodeBuilder::new("prop")
+ .attr("config_code", "100")
+ .attr("config_value", "enabled")
+ .attr("event_code", "5138")
+ .attr("sampling_weight", "-1")
+ .build();
+
+ let err = AbPropConfig::try_from_node(&node).unwrap_err();
+ assert!(err.to_string().contains("both config_code and event_code"));
+ }
+
+ #[test]
+ fn test_ab_prop_config_rejects_missing_discriminator() {
+ let node = NodeBuilder::new("prop")
+ .attr("config_value", "enabled")
+ .build();
+
+ let err = AbPropConfig::try_from_node(&node).unwrap_err();
+ assert!(err.to_string().contains("neither config_code nor event_code"));
+ }
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/src/iq/props.rs` around lines 192 - 215, Add regression tests that
exercise the two error branches in the prop parser: one where both "config_code"
and "event_code" attributes are present and one where neither is present.
Construct a node with tag "prop" and attrs containing both keys and assert that
Prop parsing returns an error mentioning both attributes (exercise the
has_config && has_event branch which currently returns Err with node.attrs);
likewise construct a node with tag "prop" and no discriminator attrs and assert
parsing returns the "neither" error (exercise the else branch). Use the same
entry point exercised by the code (the Prop parsing / TryFrom node entry such as
try_from_node_ref or the public Prop->try_from_node API) so the tests fail if
those branches regress.
Comment on lines
+129
to
+139
| /// Load multiple pre-keys by ID in a single batch operation. | ||
| /// Returns only the keys that exist. | ||
| async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> { | ||
| let mut result = Vec::with_capacity(ids.len()); | ||
| for &id in ids { | ||
| if let Some(record) = self.load_prekey(id).await? { | ||
| result.push((id, record)); | ||
| } | ||
| } | ||
| Ok(result) | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Clarify the load_prekeys_batch contract.
The default implementation preserves request order and duplicate IDs while omitting missing rows, but the doc comment only promises that existing keys are returned. Please spell out the ordering/duplicate semantics here so custom backends stay behaviorally compatible.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/src/store/traits.rs` around lines 129 - 139, Update the doc comment
for the default implementation of load_prekeys_batch to explicitly state its
contract: it preserves the input order and retains duplicates from the ids
slice, includes one entry for each id only if that id exists in storage (missing
ids are omitted), and returns a Vec<(u32, Vec<u8>)> in the same order as the
input with duplicates repeated when present; mention that custom backend
implementations should match this ordering/duplicate behavior to remain
compatible with callers relying on it.
3 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Heaptrack-guided optimizations targeting connection startup hot paths. Measured on a 1-message benchmark against the bartender mock server.
Results:
Changes:
build_iq_nodetake-by-value — eliminates deepNodetree clones through the entire IQ request pipeline (was 17K alloc calls)build_upload_prekeys_requestiterator — removes intermediateVeccollect + per-keypublic_bytes.clone()AbPropConfigdiscriminator-first parsing — checksconfig_code/event_codeattribute before parsing, avoiding double-parse wastehas_session()existence check —InMemoryBackendusescontains_key(),SqliteStoreusesSELECT EXISTSinstead of loading full session blobsload_prekeys_batch+store_prekeys_batchoverrides eliminate ~2,400 per-key#[async_trait]box allocsvalidate_digest_keybatch load — singleload_prekeys_batchcall replaces 812 individualload_prekeycallsBreaking changes
RequestUtils::build_iq_nodetakesInfoQueryby value instead of&InfoQueryPreKeyUtils::build_upload_prekeys_requesttakesimpl IntoIterator<Item = (u32, Vec<u8>)>instead of&[(u32, Vec<u8>)]Test plan
cargo clippy --all --tests— zero warningscargo test -p wacore --lib— 503 passedcargo test -p whatsapp-rust --lib— 345 passedcargo test -p whatsapp-rust-sqlite-storage --lib— 21 passedSummary by CodeRabbit
Performance
Bug Fixes