Skip to content

perf!: reduce startup allocations by ~21% and peak heap by ~31% #525

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jlucaso1 merged 1 commit into from
Apr 13, 2026
Merged

perf!: reduce startup allocations by ~21% and peak heap by ~31% #525

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions http_clients/ureq-client/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,20 +27,24 @@ impl Default for UreqHttpClient {
}

fn build_agent() -> ureq::Agent {
use ureq::config::Config;

#[allow(unused_mut)]
let mut builder = Config::builder()
// 16 KB per buffer instead of the 128 KB default.
// WA API payloads are small JSON; media uses streaming I/O.
.input_buffer_size(16 * 1024)
.output_buffer_size(16 * 1024)
.max_idle_connections(3)
.max_idle_connections_per_host(2);

#[cfg(feature = "danger-skip-tls-verify")]
{
use ureq::config::Config;
use ureq::tls::TlsConfig;
Config::builder()
.tls_config(TlsConfig::builder().disable_verification(true).build())
.build()
.into()
builder = builder.tls_config(TlsConfig::builder().disable_verification(true).build());
}

#[cfg(not(feature = "danger-skip-tls-verify"))]
{
ureq::Agent::new_with_defaults()
}
builder.build().into()
}

#[async_trait]
Expand Down
69 changes: 39 additions & 30 deletions src/prekeys.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -311,42 +311,51 @@ impl Client {
guard.backend.clone()
};

// Load each prekey referenced by the server digest and extract its public key
// Batch-load all prekeys referenced by the server digest
let loaded = match backend.load_prekeys_batch(&response.prekey_ids).await {
Ok(v) => v,
Err(e) => {
log::warn!("digestKey: failed to batch-load prekeys: {:?}, skipping", e);
return Ok(());
}
};

// Build a lookup so we preserve the server-requested order.
// Dedupe the expected count since the server may send duplicate IDs.
let loaded_map: std::collections::HashMap<u32, Vec<u8>> = loaded.into_iter().collect();
let unique_requested: std::collections::HashSet<&u32> =
response.prekey_ids.iter().collect();

if loaded_map.len() < unique_requested.len() {
log::warn!(
"digestKey: missing {} local prekeys, skipping",
unique_requested.len() - loaded_map.len()
);
return Ok(());
}

let mut prekey_pubkeys = Vec::with_capacity(response.prekey_ids.len());
for prekey_id in &response.prekey_ids {
match backend.load_prekey(*prekey_id).await {
Ok(Some(record_bytes)) => {
use prost::Message;
match waproto::whatsapp::PreKeyRecordStructure::decode(record_bytes.as_slice())
{
Ok(record) => {
if let Some(pk) = record.public_key {
prekey_pubkeys.push(pk);
} else {
log::warn!(
"digestKey: prekey {} has no public key, skipping",
prekey_id
);
return Ok(());
}
}
Err(e) => {
log::warn!(
"digestKey: failed to decode prekey {}: {}, skipping",
prekey_id,
e
);
return Ok(());
}
let Some(record_bytes) = loaded_map.get(prekey_id) else {
log::warn!("digestKey: missing local prekey {}, skipping", prekey_id);
return Ok(());
};
use prost::Message;
match waproto::whatsapp::PreKeyRecordStructure::decode(record_bytes.as_slice()) {
Ok(record) => {
if let Some(pk) = record.public_key {
prekey_pubkeys.push(pk);
} else {
log::warn!(
"digestKey: prekey {} has no public key, skipping",
prekey_id
);
return Ok(());
}
}
Ok(None) => {
log::warn!("digestKey: missing local prekey {}, skipping", prekey_id);
return Ok(());
}
Err(e) => {
log::warn!(
"digestKey: failed to load prekey {}: {:?}, skipping",
"digestKey: failed to decode prekey {}: {}, skipping",
prekey_id,
e
);
Expand Down
6 changes: 3 additions & 3 deletions src/request.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,11 +135,12 @@ impl Client {
return Err(IqError::NotConnected);
}

let default_timeout = Duration::from_secs(75);
let iq_timeout = query.timeout.unwrap_or(default_timeout);
let req_id = query
.id
.clone()
.unwrap_or_else(|| self.generate_request_id());
let default_timeout = Duration::from_secs(75);

let (tx, rx) = futures::channel::oneshot::channel();
self.response_waiters
Expand All @@ -148,7 +149,7 @@ impl Client {
.insert(req_id.clone(), tx);

let request_utils = self.get_request_utils();
let node = request_utils.build_iq_node(&query, Some(req_id.clone()));
let node = request_utils.build_iq_node(query, Some(req_id.clone()));

// Register the shutdown listener BEFORE sending to avoid a window where
// a shutdown fires between send_node() completing and listen() being called.
Expand All @@ -172,7 +173,6 @@ impl Client {

// Race the IQ response against shutdown so we fail fast on disconnect
// instead of waiting the full timeout.
let iq_timeout = query.timeout.unwrap_or(default_timeout);

futures::select! {
result = rt_timeout(&*self.runtime, iq_timeout, rx).fuse() => {
Expand Down
42 changes: 42 additions & 0 deletions storages/sqlite-storage/src/sqlite_store.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1188,6 +1188,26 @@ impl SignalStore for SqliteStore {
self.get_session_for_device(address, self.device_id).await
}

async fn has_session(&self, address: &str) -> Result<bool> {
let pool = self.pool.clone();
let device_id = self.device_id;
let address_owned = address.to_string();
self.with_semaphore(move || -> Result<bool> {
let mut conn = pool
.get()
.map_err(|e| StoreError::Connection(e.to_string()))?;
let exists = diesel::select(diesel::dsl::exists(
sessions::table
.filter(sessions::address.eq(&address_owned))
.filter(sessions::device_id.eq(device_id)),
))
.get_result(&mut conn)
.map_err(|e| StoreError::Database(e.to_string()))?;
Ok(exists)
})
.await
}

async fn put_session(&self, address: &str, session: &[u8]) -> Result<()> {
self.put_session_for_device(address, session, self.device_id)
.await
Expand Down Expand Up @@ -1346,6 +1366,28 @@ impl SignalStore for SqliteStore {
.map_err(|e| StoreError::Database(e.to_string()))?
}

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
}
Comment on lines +1369 to +1389

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ 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.


async fn remove_prekey(&self, id: u32) -> Result<()> {
let pool = self.pool.clone();
let db_semaphore = self.db_semaphore.clone();
Expand Down
11 changes: 3 additions & 8 deletions wacore/src/iq/prekeys.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -365,20 +365,15 @@ impl IqSpec for PreKeyUploadSpec {
type Response = ();

fn build_iq(&self) -> InfoQuery<'static> {
// Convert PublicKeys to 32-byte raw values for the wire
let pre_keys_bytes: Vec<(u32, Vec<u8>)> = self
.pre_keys
.iter()
.map(|(id, pk)| (*id, pk.public_key_bytes().to_vec()))
.collect();

let content = PreKeyUtils::build_upload_prekeys_request(
self.registration_id,
self.identity_key.public_key_bytes().to_vec(),
self.signed_pre_key_id,
self.signed_pre_key_public.public_key_bytes().to_vec(),
self.signed_pre_key_signature.clone(),
&pre_keys_bytes,
self.pre_keys
.iter()
.map(|(id, pk)| (*id, pk.public_key_bytes().to_vec())),
);

InfoQuery::set(
Expand Down
42 changes: 21 additions & 21 deletions wacore/src/iq/props.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ impl crate::protocol::ProtocolNode for AbProp {
}
let config_value = optional_attr(node, "config_value")
.ok_or_else(|| anyhow::anyhow!("missing config_value in prop"))?
.to_string();
.into_owned();
let config_expo_key = optional_attr(node, "config_expo_key").and_then(|s| s.parse().ok());

Ok(Self {
Expand Down Expand Up @@ -189,31 +189,31 @@ impl crate::protocol::ProtocolNode for AbPropConfig {
}

fn try_from_node_ref(node: &NodeRef<'_>) -> Result<Self, anyhow::Error> {
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
))
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Comment on lines +192 to +215

@coderabbitai coderabbitai Bot Apr 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 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.

}

let experiment_err = experiment
.err()
.unwrap_or_else(|| anyhow::anyhow!("unknown error"));
let sampling_err = sampling
.err()
.unwrap_or_else(|| anyhow::anyhow!("unknown error"));
Err(anyhow::anyhow!(
"prop did not match experiment or sampling config: experiment_err={}; sampling_err={}",
experiment_err,
sampling_err
))
}
}

Expand Down
10 changes: 5 additions & 5 deletions wacore/src/prekeys.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,17 +47,17 @@ impl PreKeyUtils {
signed_pre_key_id: u32,
signed_pre_key_public_bytes: Vec<u8>,
signed_pre_key_signature: Vec<u8>,
pre_keys: &[(u32, Vec<u8>)],
pre_keys: impl IntoIterator<Item = (u32, Vec<u8>)>,
) -> Vec<Node> {
let mut pre_key_nodes = Vec::new();
let pre_keys = pre_keys.into_iter();
let (lower, upper) = pre_keys.size_hint();
let mut pre_key_nodes = Vec::with_capacity(upper.unwrap_or(lower));
for (pre_key_id, public_bytes) in pre_keys {
let id_bytes = pre_key_id.to_be_bytes()[1..].to_vec();
let node = NodeBuilder::new("key")
.children([
NodeBuilder::new("id").bytes(id_bytes).build(),
NodeBuilder::new("value")
.bytes(public_bytes.clone())
.build(),
NodeBuilder::new("value").bytes(public_bytes).build(),
])
.build();
pre_key_nodes.push(node);
Expand Down
16 changes: 4 additions & 12 deletions wacore/src/request.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,30 +178,22 @@ impl RequestUtils {
id
}

pub fn build_iq_node(&self, query: &InfoQuery<'_>, req_id: Option<String>) -> Node {
pub fn build_iq_node(&self, query: InfoQuery<'_>, req_id: Option<String>) -> Node {
let id = req_id.unwrap_or_else(|| self.generate_request_id());

let mut builder = NodeBuilder::new("iq")
.attr("id", id)
.attr("xmlns", query.namespace)
.attr("type", query.query_type.as_str())
.attr("to", &query.to);
.attr("to", query.to);

if let Some(target) = &query.target
if let Some(target) = query.target
&& !target.is_empty()
{
builder = builder.attr("target", target);
}

if let Some(content) = &query.content {
match content {
NodeContent::Bytes(b) => builder = builder.bytes(b.clone()),
NodeContent::String(s) => builder = builder.string_content(s.clone()),
NodeContent::Nodes(n) => builder = builder.children(n.clone()),
}
}

builder.build()
builder.apply_content(query.content).build()
}

pub fn parse_iq_response(&self, response_node: &NodeRef<'_>) -> Result<(), IqError> {
Expand Down
28 changes: 28 additions & 0 deletions wacore/src/store/in_memory.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,10 @@ impl SignalStore for InMemoryBackend {
Ok(())
}

async fn has_session(&self, address: &str) -> Result<bool> {
Ok(self.state.lock().await.sessions.contains_key(address))
}

async fn delete_session(&self, address: &str) -> Result<()> {
self.state.lock().await.sessions.remove(address);
Ok(())
Expand All @@ -170,6 +174,19 @@ impl SignalStore for InMemoryBackend {
Ok(())
}

async fn store_prekeys_batch(&self, keys: &[(u32, Vec<u8>)], _uploaded: bool) -> Result<()> {
let mut state = self.state.lock().await;
for (id, record) in keys {
state.prekeys.insert(
*id,
PreKeyEntry {
record: record.clone(),
},
);
}
Ok(())
}

async fn load_prekey(&self, id: u32) -> Result<Option<Vec<u8>>> {
Ok(self
.state
Expand All @@ -180,6 +197,17 @@ impl SignalStore for InMemoryBackend {
.map(|e| e.record.clone()))
}

async fn load_prekeys_batch(&self, ids: &[u32]) -> Result<Vec<(u32, Vec<u8>)>> {
let state = self.state.lock().await;
let mut result = Vec::with_capacity(ids.len());
for &id in ids {
if let Some(entry) = state.prekeys.get(&id) {
result.push((id, entry.record.clone()));
}
}
Ok(result)
}

async fn remove_prekey(&self, id: u32) -> Result<()> {
self.state.lock().await.prekeys.remove(&id);
Ok(())
Expand Down
Loading
Loading