diff --git a/src/client/device_registry.rs b/src/client/device_registry.rs index cb44b5ea5..df04421c2 100644 --- a/src/client/device_registry.rs +++ b/src/client/device_registry.rs @@ -60,11 +60,14 @@ enum UserLookupKeys { impl UserLookupKeys { /// Returns all keys to try for lookups, in preference order. - fn all_keys(&self) -> Vec<&str> { - match self { - Self::LidWithPn { lid, pn } | Self::PnWithLid { lid, pn } => vec![lid, pn], - Self::Unknown { user } => vec![user], - } + fn all_keys(&self) -> impl Iterator { + let (first, second) = match self { + Self::LidWithPn { lid, pn } | Self::PnWithLid { lid, pn } => { + (lid.as_str(), Some(pn.as_str())) + } + Self::Unknown { user } => (user.as_str(), None), + }; + std::iter::once(first).chain(second) } /// Returns the canonical (preferred) key for storage. @@ -311,7 +314,6 @@ impl Client { self.resolve_lookup_keys(user) .await .all_keys() - .into_iter() .map(String::from) .collect() } @@ -329,19 +331,17 @@ impl Client { return true; } - // Borrowed `&str` keys (like get_devices_from_registry), bound once so both - // loops share one Vec<&str>: avoids the per-message get_lookup_keys churn. + // Borrowed keys avoid allocating the owned lookup variants on this hot path. let lookup = self.resolve_lookup_keys(user).await; - let keys = lookup.all_keys(); - for &key in &keys { + for key in lookup.all_keys() { if let Some(record) = self.device_registry_cache.get(key).await { return record.devices.iter().any(|d| d.device_id == device_id); } } let backend = self.persistence_manager.backend(); - for &key in &keys { + for key in lookup.all_keys() { match backend.get_devices(key).await { Ok(Some(record)) => { let has_device = record.devices.iter().any(|d| d.device_id == device_id); @@ -397,7 +397,6 @@ impl Client { Arc::new(record_for_cache), lookup .all_keys() - .into_iter() .chain(std::iter::once(original_user.as_str())), ) .await; @@ -464,7 +463,6 @@ impl Client { Arc::new(record_for_cache), lookup .all_keys() - .into_iter() .chain(std::iter::once(original_user.as_str())), ) .await; diff --git a/src/client/node_io.rs b/src/client/node_io.rs index 617940f1c..6d4b42d6e 100644 --- a/src/client/node_io.rs +++ b/src/client/node_io.rs @@ -159,13 +159,29 @@ impl Client { // - Critical nodes (success/failure/stream:error): inline, required for state // - Message nodes: inline, preserves arrival order for per-chat queues // (MessageHandler just enqueues + ACKs, heavy crypto runs in workers) + // - Acks/receipts: inline when unobserved; retry work detaches itself // - ib (in-band): inline, ensures offline sync tracking (expected count) // is set up before offline messages are processed // - Everything else: spawned concurrently for parallelism - let process_inline = matches!( - node.tag(), - "success" | "failure" | "stream:error" | "message" | "ib" - ); + let process_inline = match node.tag() { + "success" | "failure" | "stream:error" | "message" | "ib" => true, + // Preserve concurrent callback behavior when an app + // observes these events or every raw node. + "receipt" => { + !self.synchronous_ack + && !self.raw_node_forwarding.load(Ordering::Relaxed) + && !self.core.event_bus.has_handler_for( + wacore::types::events::EventKind::Receipt, + ) + } + "ack" => { + !self.raw_node_forwarding.load(Ordering::Relaxed) + && !self.core.event_bus.has_handler_for( + wacore::types::events::EventKind::ServerAck, + ) + } + _ => false, + }; if process_inline { self.process_decrypted_node(node).await; @@ -284,6 +300,19 @@ impl Client { self: &Arc, node: wacore_binary::OwnedNodeRef, ) { + // ACKs need shared ownership only for opt-in raw/node observers. The + // usual response-waiter path borrows the node and can skip the Arc. + if node.tag() == "ack" + && !self.raw_node_forwarding.load(Ordering::Relaxed) + && self.node_waiter_count.load(Ordering::Acquire) == 0 + && !self.offline_sync_metrics.active.load(Ordering::Acquire) + { + use wacore::xml::DisplayableNodeRef; + debug!(target: "Client/Recv", "{}", DisplayableNodeRef(node.get())); + self.handle_ack_response_inline(node.get()); + return; + } + // Wrap in Arc once - all handlers will share this same allocation let node_arc = Arc::new(node); self.process_node(node_arc).await; @@ -438,13 +467,32 @@ impl Client { return; } - // Dispatch to appropriate handler using the router - // Clone Arc (cheap - just reference count) not the Node itself - if !self - .stanza_router - .dispatch(self.clone(), Arc::clone(&node), &mut cancelled) - .await - { + // Bypass async_trait's boxed future for the hot built-in handlers while + // retaining router registration for direct router callers. + let handled = match nr.tag.as_ref() { + "ack" => { + self.handle_ack_response_inline(nr); + true + } + "receipt" => { + self.handle_receipt_inline(Arc::clone(&node)); + true + } + "message" => { + crate::handlers::message::MessageHandler::handle_inline( + self.clone(), + Arc::clone(&node), + &mut cancelled, + ) + .await + } + _ => { + self.stanza_router + .dispatch(self.clone(), Arc::clone(&node), &mut cancelled) + .await + } + }; + if !handled { warn!( "Received unknown top-level node: {}", DisplayableNodeRef(nr) @@ -1096,11 +1144,15 @@ impl Client { /// /// If an ack with an ID that matches a pending task in `response_waiters`, /// the task is resolved and the function returns `true`. Otherwise, returns `false`. + pub(crate) async fn handle_ack_response(&self, node: &wacore_binary::NodeRef<'_>) -> bool { + self.handle_ack_response_inline(node) + } + #[cfg_attr( feature = "tracing", tracing::instrument(name = "wa.conn.ack_response", level = "debug", skip_all) )] - pub(crate) async fn handle_ack_response(&self, node: &wacore_binary::NodeRef<'_>) -> bool { + pub(crate) fn handle_ack_response_inline(&self, node: &wacore_binary::NodeRef<'_>) -> bool { let ack_id = node.get_attr("id"); let ack_error = node.get_attr("error"); diff --git a/src/handlers/message.rs b/src/handlers/message.rs index 1f46338bb..597f9ae8b 100644 --- a/src/handlers/message.rs +++ b/src/handlers/message.rs @@ -15,19 +15,12 @@ const MAX_MESSAGE_DELAY_MS: u64 = 20_000; #[derive(Default)] pub struct MessageHandler; -#[cfg_attr(target_arch = "wasm32", async_trait(?Send))] -#[cfg_attr(not(target_arch = "wasm32"), async_trait)] -impl StanzaHandler for MessageHandler { - fn tag(&self) -> &'static str { - "message" - } - +impl MessageHandler { #[cfg_attr( feature = "tracing", tracing::instrument(name = "wa.recv.message_enqueue", level = "debug", skip_all) )] - async fn handle( - &self, + pub(crate) async fn handle_inline( client: Arc, node: Arc, cancelled: &mut bool, @@ -62,6 +55,23 @@ impl StanzaHandler for MessageHandler { } } +#[cfg_attr(target_arch = "wasm32", async_trait(?Send))] +#[cfg_attr(not(target_arch = "wasm32"), async_trait)] +impl StanzaHandler for MessageHandler { + fn tag(&self) -> &'static str { + "message" + } + + async fn handle( + &self, + client: Arc, + node: Arc, + cancelled: &mut bool, + ) -> bool { + Self::handle_inline(client, node, cancelled).await + } +} + /// Construct a ChatLane with a spawned worker task. Extracted to keep the /// init closure passed to `get_with_by_ref` small. fn create_chat_lane(client: &Arc) -> ChatLane { diff --git a/src/msg_secret_buffer.rs b/src/msg_secret_buffer.rs index 5e4a0b3e1..74f51d0aa 100644 --- a/src/msg_secret_buffer.rs +++ b/src/msg_secret_buffer.rs @@ -72,28 +72,41 @@ impl MsgSecretWriteBuffer { return; } { - use wacore::store::traits::{merge_msg_secret_expiry, merge_msg_secret_message_ts}; let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner()); - for mut entry in entries { - let key = ( - entry.chat.clone(), - entry.sender.clone(), - entry.msg_id.clone(), - ); - // Two captures coalescing in the same window must merge the - // retention metadata exactly like the backend upsert would - // have for two sequential writes (never-expire wins, windows - // never shrink, a known parent time is never clobbered). - if let Some(existing) = pending.get(&key) { - entry.expires_at = - merge_msg_secret_expiry(existing.expires_at, entry.expires_at); - entry.message_ts = - merge_msg_secret_message_ts(existing.message_ts, entry.message_ts); - } - pending.insert(key, entry); + for entry in entries { + Self::insert_pending(&mut pending, entry); } } - // The mutex above orders this load against seal(): an insert that the + self.schedule_or_flush().await; + } + + /// Single-entry fast path for live sends, avoiding a temporary one-element Vec. + pub(crate) async fn queue_one(self: &Arc, entry: MsgSecretEntry) { + { + let mut pending = self.pending.lock().unwrap_or_else(|p| p.into_inner()); + Self::insert_pending(&mut pending, entry); + } + self.schedule_or_flush().await; + } + + fn insert_pending(pending: &mut HashMap, mut entry: MsgSecretEntry) { + use wacore::store::traits::{merge_msg_secret_expiry, merge_msg_secret_message_ts}; + + let key = ( + entry.chat.clone(), + entry.sender.clone(), + entry.msg_id.clone(), + ); + // Coalesced captures merge retention metadata like sequential backend writes. + if let Some(existing) = pending.get(&key) { + entry.expires_at = merge_msg_secret_expiry(existing.expires_at, entry.expires_at); + entry.message_ts = merge_msg_secret_message_ts(existing.message_ts, entry.message_ts); + } + pending.insert(key, entry); + } + + async fn schedule_or_flush(self: &Arc) { + // The insertion mutex orders this load against seal(): an insert that the // shutdown flush's snapshot missed observes sealed and writes inline. if self.sealed.load(Ordering::Acquire) { self.flush().await; diff --git a/src/receipt.rs b/src/receipt.rs index 1869e810a..711a885e2 100644 --- a/src/receipt.rs +++ b/src/receipt.rs @@ -349,11 +349,15 @@ impl Client { || info.source.is_self_fanout() } + pub(crate) async fn handle_receipt(self: &Arc, node: Arc) { + self.handle_receipt_inline(node); + } + #[cfg_attr( feature = "tracing", tracing::instrument(name = "wa.receipt.handle", level = "debug", skip_all) )] - pub(crate) async fn handle_receipt(self: &Arc, node: Arc) { + pub(crate) fn handle_receipt_inline(self: &Arc, node: Arc) { let nr = node.get(); let mut attrs = nr.attrs(); let from = attrs.jid("from"); diff --git a/src/send/mod.rs b/src/send/mod.rs index 0e576c9a0..2b877d37f 100644 --- a/src/send/mod.rs +++ b/src/send/mod.rs @@ -2020,7 +2020,7 @@ impl Client { }; // Same write-behind buffer as inbound captures: visible immediately, // flushed off the send path (msmsg replies read buffer-first). - self.msg_secret_buffer.queue(vec![entry]).await; + self.msg_secret_buffer.queue_one(entry).await; } /// Decide the identity (LID vs PN) under which an outbound DM's diff --git a/wacore/src/store/in_memory.rs b/wacore/src/store/in_memory.rs index 5be65511f..9f7f2e9e4 100644 --- a/wacore/src/store/in_memory.rs +++ b/wacore/src/store/in_memory.rs @@ -5,6 +5,7 @@ //! when the struct is dropped. use std::collections::HashMap; +use std::sync::Arc; use std::sync::atomic::{AtomicI32, Ordering}; use crate::appstate::hash::HashState; @@ -146,6 +147,19 @@ impl SignalStore for InMemoryBackend { Ok(()) } + async fn put_sessions_batch(&self, sessions: &[(Arc, Bytes)]) -> Result<()> { + let mut state = self.state.lock().await; + state.sessions.reserve(sessions.len()); + for (address, session) in sessions { + if let Some(stored) = state.sessions.get_mut(address.as_ref()) { + *stored = session.clone(); + } else { + state.sessions.insert(address.to_string(), session.clone()); + } + } + Ok(()) + } + async fn has_session(&self, address: &str) -> Result { Ok(self.state.lock().await.sessions.contains_key(address)) } @@ -845,6 +859,34 @@ mod tests { is_backend::(); } + #[tokio::test] + async fn put_sessions_batch_inserts_and_updates() { + let backend = InMemoryBackend::new(); + let first: Arc = "15550000001:1@s.whatsapp.net".into(); + let second: Arc = "15550000002:2@s.whatsapp.net".into(); + + backend + .put_sessions_batch(&[ + (first.clone(), Bytes::from_static(b"first")), + (second.clone(), Bytes::from_static(b"second")), + ]) + .await + .unwrap(); + backend + .put_sessions_batch(&[(first.clone(), Bytes::from_static(b"updated"))]) + .await + .unwrap(); + + assert_eq!( + backend.get_session(&first).await.unwrap().unwrap(), + Bytes::from_static(b"updated") + ); + assert_eq!( + backend.get_session(&second).await.unwrap().unwrap(), + Bytes::from_static(b"second") + ); + } + #[tokio::test] async fn group_metadata_round_trip() { use crate::store::traits::ProtocolStore;