perf: zero-copy frame send, direct prekey encoding, covariant AttrsRef

.gitignore

@@ -10,4 +10,4 @@ docs
 .claude
 __pycache__
 .codex
-dhat-heap.json
+dhat-heap*.json

src/keepalive.rs

@@ -33,7 +33,8 @@ fn classify_keepalive_error(e: &IqError) -> KeepaliveResult {
         IqError::Socket(_)
         | IqError::Disconnected(_)
         | IqError::NotConnected
-        | IqError::InternalChannelClosed => KeepaliveResult::FatalFailure,
+        | IqError::InternalChannelClosed
+        | IqError::EncodeError(_) => KeepaliveResult::FatalFailure,
         // Exhaustive: forces a compile error when new IqError variants are added
         // so the developer must decide the classification.
         IqError::Timeout | IqError::ServerError { .. } | IqError::ParseError(_) => {

src/request.rs

@@ -24,6 +24,8 @@ pub enum IqError {
     ServerError { code: u16, text: String },
     #[error("Internal channel closed unexpectedly")]
     InternalChannelClosed,
+    #[error("Failed to encode IQ request: {0}")]
+    EncodeError(anyhow::Error),
     #[error("Failed to parse IQ response: {0}")]
     ParseError(#[from] anyhow::Error),
 }
@@ -130,39 +132,95 @@ impl Client {
         &self,
         query: InfoQuery<'_>,
     ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError> {
-        // Fail fast if the client is shutting down
-        if !self.is_running.load(Ordering::Relaxed) {
-            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 request_utils = self.get_request_utils();
+        let node = request_utils.build_iq_node(query, Some(req_id.clone()));
+
+        self.send_and_wait_iq(req_id, iq_timeout, async { self.send_node(node).await })
+            .await
+    }
+
+    /// Executes an IQ specification and returns the typed response.
+    ///
+    /// This is a convenience method that combines building the IQ request,
+    /// sending it, and parsing the response into a single operation.
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// use wacore::iq::groups::GroupQueryIq;
+    ///
+    /// let group_info = client.execute(GroupQueryIq::new(&group_jid)).await?;
+    /// println!("Group subject: {}", group_info.subject);
+    /// ```
+    pub async fn execute<S>(&self, spec: S) -> Result<S::Response, IqError>
+    where
+        S: wacore::iq::spec::IqSpec,
+    {
+        let req_id = self.generate_request_id();
+
+        // Direct-encode fast path: skip Node tree for hot IQ specs (e.g. PreKeyUploadSpec)
+        {
+            let mut buf = Vec::new();
+            match spec.encode_iq_direct(&req_id, &mut buf) {
+                Ok(true) => {
+                    let response = self
+                        .send_and_wait_iq(req_id, Duration::from_secs(75), async {
+                            self.send_raw_bytes(buf).await
+                        })
+                        .await?;
+                    return spec
+                        .parse_response(response.get())
+                        .map_err(IqError::ParseError);
+                }
+                Err(e) => return Err(IqError::EncodeError(e)),
+                Ok(false) => {}
+            }
+        }
+
+        let mut iq = spec.build_iq();
+        if iq.id.is_none() {
+            iq.id = Some(req_id);
+        }
+        let response = self.send_iq(iq).await?;
+        spec.parse_response(response.get())
+            .map_err(IqError::ParseError)
+    }
+
+    /// Centralizes waiter registration and shutdown/timeout handling.
+    async fn send_and_wait_iq<F>(
+        &self,
+        req_id: String,
+        timeout: Duration,
+        send_fn: F,
+    ) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError>
+    where
+        F: std::future::Future<Output = Result<(), crate::client::ClientError>>,
+    {
+        if !self.is_running.load(Ordering::Relaxed) {
+            return Err(IqError::NotConnected);
+        }
+
         let (tx, rx) = futures::channel::oneshot::channel();
         self.response_waiters
             .lock()
             .await
             .insert(req_id.clone(), tx);
 
-        let request_utils = self.get_request_utils();
-        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.
         let shutdown = self.shutdown_notifier.listen();
 
-        // Re-check after registering the listener to close the race window where
-        // shutdown fires between the initial check and the listen() call above.
         if !self.is_running.load(Ordering::Acquire) {
             self.response_waiters.lock().await.remove(&req_id);
             return Err(IqError::NotConnected);
         }
 
-        if let Err(e) = self.send_node(node).await {
+        if let Err(e) = send_fn.await {
             self.response_waiters.lock().await.remove(&req_id);
             return match e {
                 crate::client::ClientError::Socket(s_err) => Err(IqError::Socket(s_err)),
@@ -171,11 +229,9 @@ impl Client {
             };
         }
 
-        // Race the IQ response against shutdown so we fail fast on disconnect
-        // instead of waiting the full timeout.
-
+        let request_utils = self.get_request_utils();
         futures::select! {
-            result = rt_timeout(&*self.runtime, iq_timeout, rx).fuse() => {
+            result = rt_timeout(&*self.runtime, timeout, rx).fuse() => {
                 match result {
                     Ok(Ok(response_node)) => match request_utils.parse_iq_response(response_node.get()) {
                         Ok(()) => Ok(response_node),
@@ -194,27 +250,4 @@ impl Client {
             }
         }
     }
-
-    /// Executes an IQ specification and returns the typed response.
-    ///
-    /// This is a convenience method that combines building the IQ request,
-    /// sending it, and parsing the response into a single operation.
-    ///
-    /// # Example
-    ///
-    /// ```ignore
-    /// use wacore::iq::groups::GroupQueryIq;
-    ///
-    /// let group_info = client.execute(GroupQueryIq::new(&group_jid)).await?;
-    /// println!("Group subject: {}", group_info.subject);
-    /// ```
-    pub async fn execute<S>(&self, spec: S) -> Result<S::Response, IqError>
-    where
-        S: wacore::iq::spec::IqSpec,
-    {
-        let iq = spec.build_iq();
-        let response = self.send_iq(iq).await?;
-        spec.parse_response(response.get())
-            .map_err(IqError::ParseError)
-    }
 }

src/socket/noise_socket.rs

@@ -78,9 +78,10 @@ impl NoiseSocket {
         send_job_rx: async_channel::Receiver<SendJob>,
     ) {
         let mut write_counter: u32 = 0;
-        // Reusable buffers -- capacity stays allocated between sends
         let mut enc_buf = Vec::with_capacity(4096);
-        let mut out_buf = Vec::with_capacity(4096);
+        // BytesMut: split().freeze() yields a zero-copy Bytes while retaining
+        // the underlying allocation for the next frame.
+        let mut out_buf = BytesMut::with_capacity(4096);
 
         while let Ok(job) = send_job_rx.recv().await {
             let result = Self::process_send_job(
@@ -106,24 +107,21 @@ impl NoiseSocket {
         write_counter: &mut u32,
         plaintext: &[u8],
         enc_buf: &mut Vec<u8>,
-        out_buf: &mut Vec<u8>,
+        out_buf: &mut BytesMut,
     ) -> SendResult {
         let counter = *write_counter;
 
         if plaintext.len() <= INLINE_ENCRYPT_THRESHOLD {
-            // Copy into reusable enc_buf, encrypt in place
             enc_buf.clear();
             enc_buf.extend_from_slice(plaintext);
             if let Err(e) = write_key.encrypt_in_place_with_counter(counter, enc_buf) {
                 return Err(EncryptSendError::crypto(anyhow::anyhow!(e.to_string())));
             }
 
-            out_buf.clear();
             if let Err(e) = wacore::framing::encode_frame_into(enc_buf, None, out_buf) {
                 return Err(EncryptSendError::framing(e));
             }
         } else {
-            // Large messages: encrypt on blocking thread (reads plaintext, returns new ciphertext)
             let write_key = write_key.clone();
             let plaintext_owned = plaintext.to_vec();
 
@@ -139,15 +137,15 @@ impl NoiseSocket {
                 }
             };
 
-            out_buf.clear();
             if let Err(e) = wacore::framing::encode_frame_into(&ciphertext, None, out_buf) {
                 return Err(EncryptSendError::framing(e));
             }
         }
 
-        // copy_from_slice so out_buf retains its capacity for the next send
-        let frame = bytes::Bytes::copy_from_slice(out_buf);
-        out_buf.clear();
+        // Zero-copy: split() moves the written data into a new BytesMut,
+        // freeze() converts it to Bytes. The original out_buf retains its
+        // allocated capacity for the next frame.
+        let frame = out_buf.split().freeze();
         if let Err(e) = transport.send(frame).await {
             return Err(EncryptSendError::transport(e));
         }

wacore/binary/src/decoder.rs

@@ -425,17 +425,20 @@ impl<'a> Decoder<'a> {
     }
 
     fn read_attributes(&mut self, size: usize) -> Result<AttrsRef<'a>> {
-        let mut attrs = AttrsRef::with_capacity(size);
+        if size == 0 {
+            return Ok(AttrsRef::Empty);
+        }
+        let mut v = Vec::with_capacity(size);
         for _ in 0..size {
             let key = self
                 .read_value_as_string()?
                 .ok_or(BinaryError::NonStringKey)?;
             let value = self
                 .read_value()?
                 .unwrap_or(ValueRef::String(NodeStr::Borrowed("")));
-            attrs.push((key, value));
+            v.push((key, value));
         }
-        Ok(attrs)
+        Ok(AttrsRef::from_vec(v))
     }
 
     fn read_content(&mut self) -> Result<Option<NodeContentRef<'a>>> {

wacore/binary/src/encoder.rs

@@ -12,7 +12,7 @@ use crate::jid::{self, Jid, JidRef};
 use crate::node::{Node, NodeContent, NodeContentRef, NodeRef, NodeValue, ValueRef};
 use crate::token;
 
-pub(crate) trait ByteWriter {
+pub trait ByteWriter {
     fn write_u8(&mut self, value: u8) -> Result<()>;
     fn write_bytes(&mut self, bytes: &[u8]) -> Result<()>;
 }
@@ -41,7 +41,7 @@ impl<W: Write> ByteWriter for IoByteWriter<W> {
     }
 }
 
-pub(crate) struct VecByteWriter<'a> {
+pub struct VecByteWriter<'a> {
     buffer: &'a mut Vec<u8>,
 }
 
@@ -110,7 +110,7 @@ impl ByteWriter for SliceByteWriter<'_> {
 /// Trait for encoding node structures (both owned Node and borrowed NodeRef).
 /// All encoding logic lives in the trait implementation, keeping
 /// the Encoder simple and focused on low-level byte writing.
-pub(crate) trait EncodeNode {
+pub trait EncodeNode {
     fn tag(&self) -> &str;
     fn attrs_len(&self) -> usize;
     fn has_content(&self) -> bool;
@@ -177,7 +177,7 @@ impl EncodeNode for NodeRef<'_> {
     }
 
     fn encode_attrs<'a, W: ByteWriter>(&self, encoder: &mut Encoder<'a, W>) -> Result<()> {
-        for (k, v) in &self.attrs {
+        for (k, v) in self.attrs.iter() {
             encoder.write_string(k)?;
             match v {
                 ValueRef::String(s) => encoder.write_string(s)?,
@@ -602,13 +602,13 @@ fn validate_hex(value: &str) -> bool {
         .all(|&b| b.is_ascii_digit() || (b'A'..=b'F').contains(&b))
 }
 
-pub(crate) struct Encoder<'a, W: ByteWriter> {
+pub struct Encoder<'a, W: ByteWriter> {
     writer: W,
     string_hints: Option<&'a StringHintCache>,
 }
 
 impl<W: Write> Encoder<'static, IoByteWriter<W>> {
-    pub(crate) fn new(writer: W) -> Result<Self> {
+    pub fn new(writer: W) -> Result<Self> {
         let mut enc = Self {
             writer: IoByteWriter::new(writer),
             string_hints: None,
@@ -619,7 +619,8 @@ impl<W: Write> Encoder<'static, IoByteWriter<W>> {
 }
 
 impl<'v> Encoder<'static, VecByteWriter<'v>> {
-    pub(crate) fn new_vec(buffer: &'v mut Vec<u8>) -> Result<Self> {
+    pub fn new_vec(buffer: &'v mut Vec<u8>) -> Result<Self> {
+        buffer.clear();
         let mut enc = Self {
             writer: VecByteWriter::new(buffer),
             string_hints: None,
@@ -680,7 +681,7 @@ impl<'a, W: ByteWriter> Encoder<'a, W> {
     }
 
     #[inline(always)]
-    fn write_bytes_with_len(&mut self, bytes: &[u8]) -> Result<()> {
+    pub fn write_bytes_with_len(&mut self, bytes: &[u8]) -> Result<()> {
         let len = bytes.len();
         if len < 256 {
             self.write_u8(token::BINARY_8)?;
@@ -696,7 +697,7 @@ impl<'a, W: ByteWriter> Encoder<'a, W> {
     }
 
     #[inline(always)]
-    fn write_string(&mut self, s: &str) -> Result<()> {
+    pub fn write_string(&mut self, s: &str) -> Result<()> {
         if let Some(string_hints) = self.string_hints
             && let Some(hint) = string_hints.hint_for(s)
         {
@@ -782,7 +783,7 @@ impl<'a, W: ByteWriter> Encoder<'a, W> {
 
     /// Write an owned Jid directly without converting to string first.
     /// This avoids the allocation that would occur with `jid.to_string()`.
-    fn write_jid_owned(&mut self, jid: &Jid) -> Result<()> {
+    pub fn write_jid_owned(&mut self, jid: &Jid) -> Result<()> {
         if jid.device > 0 {
             // AD_JID format: domain_type, device, user
             let device = u8::try_from(jid.device).map_err(|_| {
@@ -909,21 +910,23 @@ impl<'a, W: ByteWriter> Encoder<'a, W> {
         Ok(())
     }
 
-    fn write_list_start(&mut self, len: usize) -> Result<()> {
+    pub fn write_list_start(&mut self, len: usize) -> Result<()> {
         if len == 0 {
             self.write_u8(token::LIST_EMPTY)?;
         } else if len < 256 {
             self.write_u8(248)?;
             self.write_u8(len as u8)?;
-        } else {
+        } else if len <= u16::MAX as usize {
             self.write_u8(249)?;
             self.write_u16_be(len as u16)?;
+        } else {
+            return Err(BinaryError::InvalidNode);
         }
         Ok(())
     }
 
     /// Write any node type (owned or borrowed) using the EncodeNode trait.
-    pub(crate) fn write_node<N: EncodeNode>(&mut self, node: &N) -> Result<()> {
+    pub fn write_node<N: EncodeNode>(&mut self, node: &N) -> Result<()> {
         let content_len = if node.has_content() { 1 } else { 0 };
         let list_len = 1 + (node.attrs_len() * 2) + content_len;
 

wacore/binary/src/lib.rs

@@ -4,7 +4,7 @@ pub mod attrs;
 pub mod builder;
 pub mod consts;
 mod decoder;
-mod encoder;
+pub mod encoder;
 pub mod error;
 pub mod jid;
 pub mod marshal;