perf: zero-copy frame send, direct prekey encoding, covariant AttrsRef - #552
Conversation
Make encode_frame_into generic over a FrameBuf trait (implemented for both Vec<u8> and BytesMut). Change the NoiseSocket sender_task's out_buf from Vec<u8> to BytesMut. split().freeze() yields a Bytes view of the written frame data while the BytesMut retains its underlying allocation for the next frame. This eliminates the per-frame Bytes::copy_from_slice that was copying ~1KB per outbound message.
Add encode_iq_direct fast path to IqSpec trait. PreKeyUploadSpec implements it with a custom EncodeNode that writes the entire <iq> stanza directly via the Encoder, bypassing NodeBuilder and marshal. This eliminates 812 Node object allocations (4,872 blocks) and the marshal_auto serialization pass for the prekey upload IQ. The encoder writes tags, attributes, and byte content directly into a pre-sized Vec without building an intermediate Node tree. Made Encoder, EncodeNode, ByteWriter, and key write methods public so crates outside wacore-binary can implement custom encoding.
… Vec) Replace AttrsRef type alias (Vec) with a covariant enum that stores 0-1 attributes inline without heap allocation. Most WA protocol nodes have 3-5 attributes so they still use the Many(Vec) variant, but zero-attr and single-attr nodes (common in child elements) avoid the Vec entirely. The enum is covariant in the lifetime parameter (unlike SmallVec), so it works with yoke::Yokeable for the zero-copy decode path. Manual Yokeable impl uses transmute for lifetime erasure. Implements: len, is_empty, iter, as_slice, push, with_capacity, FromIterator, and serde::Serialize (via as_slice delegation).
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughClient gained a direct-binary IQ encode fast path and a centralized send-and-wait helper; encoder APIs and encoder module were made public; frame buffering generalized to support Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Spec as IqSpec
participant Encoder
participant Socket
participant Waiter as ResponseWaiter
Note over Client,Waiter: Direct-encode IQ flow (new)
Client->>Spec: encode_iq_direct(req_id, buf)
alt Ok(true)
Spec->>Encoder: Encoder::new_vec(buf)
Encoder-->>Spec: writes encoded bytes
Spec-->>Client: Ok(true)
Client->>Waiter: register waiter keyed by req_id
Client->>Socket: send pre-encoded bytes via send_and_wait_iq(timeout)
alt send ok
Waiter-->>Client: response / timeout / shutdown
Client->>Spec: parse_response(node)
Spec-->>Client: parsed Response or ParseError
else send err
Client->>Waiter: remove waiter
Client-->>Client: return mapped IqError
end
else Ok(false) or Err
Spec-->>Client: Ok(false) / Err
Client->>Spec: build_iq() + marshal (fallback)
Client->>Socket: send marshaled IQ via existing send path
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 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 |
Benchmark Results4 improvement(s):
55 unchanged benchmark(s)
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d830750e03
ℹ️ 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".
| match n { | ||
| 0 => Self::Empty, | ||
| _ => Self::Many(Vec::with_capacity(n)), |
There was a problem hiding this comment.
Preserve one-attribute fast path in AttrsRef::with_capacity
with_capacity(1) currently returns Many(Vec::with_capacity(1)), so callers that know there is exactly one attribute (notably Decoder::read_attributes(size)) still pay a heap allocation on the first push. That bypasses the new inline One variant for a common case (<count value="..."/>, similar single-attr nodes) and negates part of the intended allocation reduction from this commit. Initializing n == 1 as Empty (or otherwise routing first insert to One) keeps the single-attribute decode path allocation-free.
Useful? React with 👍 / 👎.
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 (3)
wacore/binary/src/encoder.rs (2)
621-628:⚠️ Potential issue | 🟠 MajorClear the
Vecbefore starting a new encoded stanza.
new_vec()currently appends the leading format byte to whatever is already inbuffer. That breaks the new reusable-buffer use case: calling it with a non-emptyVec<u8>will prefix stale bytes and produce a malformed frame.Suggested fix
impl<'v> Encoder<'static, VecByteWriter<'v>> { pub fn new_vec(buffer: &'v mut Vec<u8>) -> Result<Self> { + buffer.clear(); let mut enc = Self { writer: VecByteWriter::new(buffer), string_hints: None, }; enc.write_u8(0)?; Ok(enc) } }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/binary/src/encoder.rs` around lines 621 - 628, The Encoder::new_vec currently appends the leading format byte to an existing buffer; modify Encoder::new_vec so it clears the provided Vec<u8> before constructing the VecByteWriter (or before writing the leading byte) to ensure the buffer is empty for a new stanza; locate the new_vec function in impl<'v> Encoder<'static, VecByteWriter<'v>> and call clear on the passed-in buffer (or ensure VecByteWriter is initialized with an emptied buffer) prior to enc.write_u8(0).
912-920:⚠️ Potential issue | 🟠 MajorReject list lengths that exceed
LIST_16capacity.This now-public API silently truncates
lenwithlen as u16for values above65535, which emits corrupt wire data instead of failing fast. Please return an error when the length is out of range.Suggested fix
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 if u16::try_from(len).is_ok() { + let len = len as u16; + self.write_u8(249)?; + self.write_u16_be(len)?; } else { - self.write_u8(249)?; - self.write_u16_be(len as u16)?; + return Err(BinaryError::AttrParse(format!( + "list length exceeds LIST_16 max: {}", + len + ))); } Ok(()) }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/binary/src/encoder.rs` around lines 912 - 920, The write_list_start function currently truncates oversized lengths when taking the LIST_16 path (it casts len to u16 before calling write_u16_be), so update write_list_start to explicitly reject lengths > 65535 by returning an error (Result::Err) instead of silently truncating; in practice add a check before the else branch (or inside the else branch for the 249 case) that if len > u16::MAX returns an appropriate error (e.g., invalid input/length too large) so write_u16_be is only called with a safe u16 conversion.wacore/src/iq/prekeys.rs (1)
365-490: 🧹 Nitpick | 🔵 TrivialAdd a byte-for-byte equivalence test for the direct encoder.
This path now hand-writes the entire
<iq>stanza, but the tests in this file still only coverbuild_iq()and response parsing. Please add a regression test thatencode_iq_direct()produces the same bytes asbuild_iq()+ marshal for the samePreKeyUploadSpec, so future protocol drift is caught locally instead of against the server.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/src/iq/prekeys.rs` around lines 365 - 490, Add a unit test that constructs a representative PreKeyUploadSpec, then: 1) call PreKeyUploadSpec::encode_iq_direct(request_id, &mut Vec) (which uses PreKeyUploadIqNode) to produce bytes, and 2) call PreKeyUploadSpec::build_iq() to get the InfoQuery/Node and marshal it with the existing node marshaller to produce bytes; finally assert the two byte buffers are identical. Place the test alongside the existing tests in this file, use the same request_id and spec data for both paths, and fail on any mismatch to catch protocol drift.
🤖 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/binary/src/node.rs`:
- Around line 398-401: with_capacity currently maps n == 1 into
Many(Vec::with_capacity(1)) which forces a heap allocation and bypasses the
inline One variant; change AttrsRef::with_capacity so it explicitly returns the
One variant when n == 1 (construct the appropriate single-attribute placeholder
that One expects) and only return Many(Vec::with_capacity(n)) for n > 1, keeping
0 => Empty; update the with_capacity function and any constructors used by
AttrsRef::One to ensure the One variant is created correctly without allocating.
In `@wacore/src/iq/spec.rs`:
- Around line 14-24: The fast-path method encode_iq_direct currently only
returns a bool and drops InfoQuery.timeout so Client::execute cannot preserve
non-default timeouts; change the contract to surface timeout metadata (e.g.
change encode_iq_direct to return Result<(bool, Option<std::time::Duration>),
anyhow::Error> or accept a &mut Option<Duration> timeout out-parameter), update
all implementations to populate the timeout from InfoQuery.timeout when
returning Ok(true), and update Client::execute to consume that timeout instead
of falling back to the hard-coded 75s when the fast path is used; ensure
build_iq and any callers still work with the revised signature.
---
Outside diff comments:
In `@wacore/binary/src/encoder.rs`:
- Around line 621-628: The Encoder::new_vec currently appends the leading format
byte to an existing buffer; modify Encoder::new_vec so it clears the provided
Vec<u8> before constructing the VecByteWriter (or before writing the leading
byte) to ensure the buffer is empty for a new stanza; locate the new_vec
function in impl<'v> Encoder<'static, VecByteWriter<'v>> and call clear on the
passed-in buffer (or ensure VecByteWriter is initialized with an emptied buffer)
prior to enc.write_u8(0).
- Around line 912-920: The write_list_start function currently truncates
oversized lengths when taking the LIST_16 path (it casts len to u16 before
calling write_u16_be), so update write_list_start to explicitly reject lengths >
65535 by returning an error (Result::Err) instead of silently truncating; in
practice add a check before the else branch (or inside the else branch for the
249 case) that if len > u16::MAX returns an appropriate error (e.g., invalid
input/length too large) so write_u16_be is only called with a safe u16
conversion.
In `@wacore/src/iq/prekeys.rs`:
- Around line 365-490: Add a unit test that constructs a representative
PreKeyUploadSpec, then: 1) call PreKeyUploadSpec::encode_iq_direct(request_id,
&mut Vec) (which uses PreKeyUploadIqNode) to produce bytes, and 2) call
PreKeyUploadSpec::build_iq() to get the InfoQuery/Node and marshal it with the
existing node marshaller to produce bytes; finally assert the two byte buffers
are identical. Place the test alongside the existing tests in this file, use the
same request_id and spec data for both paths, and fail on any mismatch to catch
protocol drift.
🪄 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: ba178aaf-a2df-4ea1-8e59-bf9eba48b2ea
📒 Files selected for processing (8)
src/request.rssrc/socket/noise_socket.rswacore/binary/src/encoder.rswacore/binary/src/lib.rswacore/binary/src/node.rswacore/noise/src/framing.rswacore/src/iq/prekeys.rswacore/src/iq/spec.rs
| pub fn with_capacity(n: usize) -> Self { | ||
| match n { | ||
| 0 => Self::Empty, | ||
| _ => Self::Many(Vec::with_capacity(n)), |
There was a problem hiding this comment.
with_capacity(1) still allocates and bypasses One.
decoder.read_attributes() builds through AttrsRef::with_capacity(size), so the size == 1 case takes this Many(Vec::with_capacity(1)) branch and never uses the inline One variant. That means the hot single-attribute decode path still heap-allocates, which defeats the main optimization this enum is supposed to provide.
Suggested fix
pub fn with_capacity(n: usize) -> Self {
match n {
- 0 => Self::Empty,
+ 0 | 1 => Self::Empty,
_ => Self::Many(Vec::with_capacity(n)),
}
}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/binary/src/node.rs` around lines 398 - 401, with_capacity currently
maps n == 1 into Many(Vec::with_capacity(1)) which forces a heap allocation and
bypasses the inline One variant; change AttrsRef::with_capacity so it explicitly
returns the One variant when n == 1 (construct the appropriate single-attribute
placeholder that One expects) and only return Many(Vec::with_capacity(n)) for n
> 1, keeping 0 => Empty; update the with_capacity function and any constructors
used by AttrsRef::One to ensure the One variant is created correctly without
allocating.
| /// Optionally encode the IQ stanza directly into a pre-sized buffer, | ||
| /// bypassing the Node intermediate representation. Returns `true` if | ||
| /// the fast path was used; `false` falls back to `build_iq()` + marshal. | ||
| /// | ||
| /// The buffer must contain the full binary-encoded `<iq>` stanza including | ||
| /// the leading format byte. `request_id` is the IQ request ID. | ||
| fn encode_iq_direct( | ||
| &self, | ||
| _request_id: &str, | ||
| _out: &mut Vec<u8>, | ||
| ) -> Result<bool, anyhow::Error> { |
There was a problem hiding this comment.
The fast-path API drops IQ timeout metadata.
build_iq() can carry a non-default InfoQuery.timeout, but encode_iq_direct() only hands back bytes. Once a spec returns Ok(true), Client::execute has no way to preserve that timeout and falls back to the hard-coded 75s path instead. Please surface timeout metadata through the fast-path contract as well, or the optimized path will change request behavior for those specs.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/src/iq/spec.rs` around lines 14 - 24, The fast-path method
encode_iq_direct currently only returns a bool and drops InfoQuery.timeout so
Client::execute cannot preserve non-default timeouts; change the contract to
surface timeout metadata (e.g. change encode_iq_direct to return Result<(bool,
Option<std::time::Duration>), anyhow::Error> or accept a &mut Option<Duration>
timeout out-parameter), update all implementations to populate the timeout from
InfoQuery.timeout when returning Ok(true), and update Client::execute to consume
that timeout instead of falling back to the hard-coded 75s when the fast path is
used; ensure build_iq and any callers still work with the revised signature.
|
@claude review this nitpicky check claims, DRY usage, unnecessary allocations, better options to the same solution, if it lacks in tests, if it can have something better to reuse (like if we want to reuse the same pattern for another). |
|
Claude finished @jlucaso1's task in 4m 22s —— View job Review of PR #552I've read every changed file and cross-referenced against the existing codebase. Here's a detailed review covering the specific areas you asked about. 1. Nitpicky Reviewer Claims — VerifiedCodeRabbit: " CodeRabbit: " CodeRabbit: " Codex/CodeRabbit: " 2. DRY Analysis
The only differences are: (a) Suggestion: Extract a shared async fn send_and_wait_iq<F, Fut>(
&self,
req_id: String,
timeout: Duration,
send_fn: F,
) -> Result<Arc<wacore_binary::OwnedNodeRef>, IqError>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<(), crate::client::ClientError>>,
{ /* shared logic */ }3. Unnecessary Allocations
4. Better Alternatives
Enum size consideration for 5. Test Coverage GapsNo equivalence test between #[test]
fn test_direct_encode_matches_build_iq() {
let spec = PreKeyUploadSpec::new(/* ... */);
let req_id = "test-123";
// Fast path
let mut direct_buf = Vec::new();
spec.encode_iq_direct(req_id, &mut direct_buf).unwrap();
// Slow path
let iq = spec.build_iq();
let request_utils = /* ... */;
let node = request_utils.build_iq_node(iq, Some(req_id.to_string()));
let marshal_buf = wacore_binary::marshal(&node).unwrap();
assert_eq!(direct_buf, marshal_buf);
}No No test for No unit test for 6. Reusability — Can Another Spec Use This Pattern?The
For For most other specs, it's not worth it. The ~50 IQ specs in If you want reusability, consider a middle ground: a macro or builder that generates the Summary of Actionable Items
| |
- AttrsRef::with_capacity(1) now returns Empty instead of Many(Vec), fixing the benchmark regressions (+10-17% on unmarshal/roundtrip) caused by heap-allocating for single-attr nodes - Encoder::new_vec clears buffer before writing format byte - write_list_start rejects len > u16::MAX instead of truncating - Vec::new() in execute() moved inside the fast-path block to avoid allocation when encode_iq_direct returns false (all non-prekey specs) - Added unit test verifying encode_iq_direct produces identical bytes to build_iq + marshal
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/request.rs`:
- Line 244: The hardcoded default_timeout (Duration::from_secs(75)) reduces
flexibility for different IqSpec implementations; update the API so callers can
control timeout instead of using a fixed 75s. Specifically, modify
encode_iq_direct (or its return type) or add an optional trait method on IqSpec
to provide a per-spec timeout, then replace the direct use of default_timeout
with the provided timeout (falling back to a sensible default) and update
send_iq to accept or honor that timeout; reference symbols: default_timeout,
Duration::from_secs(75), encode_iq_direct, send_iq, IqSpec.
- Around line 233-288: The send_iq_raw implementation duplicates logic from
send_iq; extract a shared helper (e.g., handle_iq_request or
register_waiter_and_send) that encapsulates: checking is_running (use same
ordering semantics), registering/removing a oneshot waiter in response_waiters,
creating shutdown listener from shutdown_notifier, invoking send_raw_bytes (or
send_bytes closure passed in), running rt_timeout with the runtime, parsing
response via get_request_utils().parse_iq_response, and cleaning up the waiter
on timeout/shutdown/errors; then refactor send_iq_raw and send_iq to call this
helper with the appropriate request id, buffer/encode step, and error mapping so
the duplicated blocks (is_running checks, waiter registration/removal,
shutdown/timeout race, and response parsing) are centralized.
In `@wacore/binary/src/node.rs`:
- Around line 464-485: The unsafe Yokeable impl for AttrsRef (impl Yokeable for
AttrsRef<'static>) relies on AttrsRef remaining covariant in 'a; add a
compile-time covariance assertion to detect future regressions by introducing a
private phantom-check function that references AttrsRef with contravariant and
covariant contexts (e.g., a helper that accepts fn(AttrsRef<'static>) ->
AttrsRef<'_> or uses core::mem::transmute checks) and call it in the module so
the compiler will error if covariance is violated; update or document this check
alongside the unsafe methods make() and transform_mut() so future changes to
AttrsRef variants (Empty, One, Many) are caught early.
🪄 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: a015d5f0-49b6-498d-adf7-e9bf8f64dcc2
📒 Files selected for processing (4)
src/request.rswacore/binary/src/encoder.rswacore/binary/src/node.rswacore/src/iq/prekeys.rs
| return Err(IqError::NotConnected); | ||
| } | ||
|
|
||
| let default_timeout = Duration::from_secs(75); |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Hardcoded timeout limits flexibility for future specs.
The 75-second default matches send_iq, but if future IqSpec implementations using encode_iq_direct need different timeouts, they'll be stuck with this value. Consider exposing an optional timeout in encode_iq_direct's return type or as a separate trait method if this becomes a practical concern.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/request.rs` at line 244, The hardcoded default_timeout
(Duration::from_secs(75)) reduces flexibility for different IqSpec
implementations; update the API so callers can control timeout instead of using
a fixed 75s. Specifically, modify encode_iq_direct (or its return type) or add
an optional trait method on IqSpec to provide a per-spec timeout, then replace
the direct use of default_timeout with the provided timeout (falling back to a
sensible default) and update send_iq to accept or honor that timeout; reference
symbols: default_timeout, Duration::from_secs(75), encode_iq_direct, send_iq,
IqSpec.
| unsafe impl<'a> yoke::Yokeable<'a> for AttrsRef<'static> { | ||
| type Output = AttrsRef<'a>; | ||
|
|
||
| fn transform(&'a self) -> &'a Self::Output { | ||
| self | ||
| } | ||
|
|
||
| fn transform_owned(self) -> Self::Output { | ||
| self | ||
| } | ||
|
|
||
| unsafe fn make(from: Self::Output) -> Self { | ||
| unsafe { std::mem::transmute(from) } | ||
| } | ||
|
|
||
| fn transform_mut<F>(&'a mut self, f: F) | ||
| where | ||
| F: 'static + for<'b> FnOnce(&'b mut Self::Output), | ||
| { | ||
| unsafe { f(std::mem::transmute::<&mut Self, &mut Self::Output>(self)) } | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Unsafe Yokeable impl looks correct but warrants careful maintenance.
The safety argument relies on AttrsRef being covariant in 'a. This holds because:
Emptyhas no lifetimeOnecontains(NodeStr<'a>, ValueRef<'a>)which are both#[derive(Yokeable)]Many(Vec<T>)is covariant inT
The transmute in make() and transform_mut() is safe given identical memory layout across lifetimes. Consider adding a compile-time covariance assertion (e.g., a phantom function) to catch accidental future changes that break covariance.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/binary/src/node.rs` around lines 464 - 485, The unsafe Yokeable impl
for AttrsRef (impl Yokeable for AttrsRef<'static>) relies on AttrsRef remaining
covariant in 'a; add a compile-time covariance assertion to detect future
regressions by introducing a private phantom-check function that references
AttrsRef with contravariant and covariant contexts (e.g., a helper that accepts
fn(AttrsRef<'static>) -> AttrsRef<'_> or uses core::mem::transmute checks) and
call it in the module so the compiler will error if covariance is violated;
update or document this check alongside the unsafe methods make() and
transform_mut() so future changes to AttrsRef variants (Empty, One, Many) are
caught early.
- DRY: extract shared send_and_wait_iq() that both send_iq() and the direct-encode fast path use. Eliminates the duplicated send_iq_raw method (~55 lines of identical waiter/shutdown/timeout logic). - Add compile-time covariance assertion for AttrsRef: if a future change makes AttrsRef invariant (breaking yoke::Yokeable safety), _assert_attrs_ref_covariant will fail to compile.
There was a problem hiding this comment.
Actionable comments posted: 5
🤖 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/request.rs`:
- Line 186: Trim the doc comment above the shared helper that “register[s]
waiter, send[s], race[s] response vs shutdown/timeout” to a one-line rationale
explaining why the helper exists (e.g., to centralize waiter registration and
consistent shutdown/timeout handling to avoid duplication and ensure uniform
cancellation semantics), removing the mechanical restatement of what it does;
update the comment that sits immediately above that helper in src/request.rs to
focus on this rationale.
- Around line 164-181: The code pre-generates req_id via generate_request_id()
and uses it on the direct-encode fast path, but on the fallback path it calls
spec.build_iq() which may leave the IQ id empty so send_iq() will generate a new
id; change execute() so after let iq = spec.build_iq(); if iq.id is empty/None
assign the previously created req_id (the one from generate_request_id()) to iq
before calling send_iq(), ensuring send_iq/send_and_wait_iq use the same request
id and avoiding the extra allocation; reference functions/fields:
generate_request_id, encode_iq_direct, build_iq, send_iq, send_and_wait_iq,
req_id.
- Around line 168-177: The current if let Ok(true) = spec.encode_iq_direct(...)
silently treats Err(_) as "not implemented" and falls back; replace that
conditional with a match on spec.encode_iq_direct(&req_id, &mut buf) so that
Ok(true) executes the fast-path (calling send_and_wait_iq/send_raw_bytes and
parse_response), Ok(false) continues to the existing marshal path, and Err(e) is
returned (or mapped) instead of being suppressed — e.g., return
Err(IqError::EncodeError(e)) (or propagate the underlying error) from the code
that currently calls encode_iq_direct in src/request.rs so failures in
encode_iq_direct are visible to client.execute().
In `@wacore/binary/src/node.rs`:
- Around line 390-395: The public enum AttrsRef with variants Empty/One/Many
allows callers to build Many(vec![]) or Many(vec![...]) making is_empty() and
derived PartialEq incorrect; change AttrsRef to a pub struct wrapper that
contains a private enum (or make the enum itself private) and expose controlled
constructors (e.g., AttrsRef::empty(), ::one(key,val), ::many(vec)) that
normalize inputs (convert empty vec -> Empty, single-element vec -> One) and
implement is_empty() and PartialEq manually (or derive on the normalized
internal state) so logically equivalent representations compare equal and empty
detection is correct; update any usages that construct AttrsRef directly to use
the new constructors (refer to AttrsRef, Empty/One/Many, is_empty, and
PartialEq).
- Around line 449-456: The FromIterator implementation for AttrsRef currently
always starts from AttrsRef::Empty and repeatedly grows, reintroducing
reallocations; fix by calling let mut iter = iter.into_iter(); use
iter.size_hint() to detect an exact/upper bound (e.g. let (_, high) =
iter.size_hint(); if let Some(sz) = high { result = AttrsRef::with_capacity(sz);
}) before the for loop, then push each item as before; this preserves
one-allocation behavior for exact-size iterators (see FromIterator for AttrsRef,
AttrsRef::Empty, result.push and Node::as_node_ref()).
🪄 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: f64477a0-3d68-400f-8e90-3519024871f3
📒 Files selected for processing (2)
src/request.rswacore/binary/src/node.rs
- execute() now propagates Err from encode_iq_direct instead of silently falling back to build_iq - Reuse the pre-generated req_id on the fallback path (set iq.id if empty) to avoid a redundant generate_request_id call - FromIterator for AttrsRef uses size_hint to pre-allocate with with_capacity for exact-size iterators
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/request.rs`:
- Around line 169-180: The outbound encode failure from spec.encode_iq_direct is
incorrectly mapped to IqError::ParseError; change the Err(e) arm so encoding
failures are returned as a distinct error (e.g., IqError::EncodeError or
propagate the underlying error) instead of ParseError, and update the IqError
enum/type if necessary; keep the Ok(true) fast-path behavior (send_and_wait_iq
-> parse_response) unchanged and only alter the Err(e) branch to return the
proper encode-level error.
🪄 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: dffd0005-e632-499a-853a-9d13dd66ea59
📒 Files selected for processing (2)
src/request.rswacore/binary/src/node.rs
…ession The One variant stored a full (NodeStr, ValueRef) tuple inline (~72 bytes), inflating every NodeRef by ~56 bytes and causing +9-17% instruction count regressions on small-node benchmarks. Replace with Empty | Slice(Box<[T]>): Box<[T]> is 16 bytes (ptr + len), making the enum 24 bytes -- same as the original Vec<T>. Zero-attr nodes still skip allocation via Empty. The decoder builds a Vec with exact capacity then converts to boxed slice via into_boxed_slice(). Also fixes: clippy mut warning, propagates encode_iq_direct errors, reuses req_id on fallback path.
Encoding failures from the direct-encode fast path were incorrectly mapped to IqError::ParseError. Added distinct EncodeError variant and classified it as FatalFailure in keepalive error handling.
Reserve capacity for the estimated wire size (812 prekeys * 40 bytes + 256 bytes overhead) before encoding. Avoids incremental Vec growth during the encode pass. Marginal impact (~121 KB, -0.4%) since the Vec growth strategy already amortizes well.
There was a problem hiding this comment.
Actionable comments posted: 3
♻️ Duplicate comments (1)
wacore/binary/src/node.rs (1)
393-420:⚠️ Potential issue | 🟠 MajorNormalize empty
AttrsRefsemantics or hide the variants.Because
AttrsRefis public, callers can still constructAttrsRef::Slice(Box::new([])). With derivedPartialEqandis_empty()matching onlyEmpty, that logically empty value reports non-empty and compares unequal toAttrsRef::Empty.Suggested fix
-#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum AttrsRef<'a> { Empty, Slice(Box<[(NodeStr<'a>, ValueRef<'a>)]>), } impl<'a> AttrsRef<'a> { @@ #[inline] pub fn is_empty(&self) -> bool { - matches!(self, Self::Empty) + self.as_slice().is_empty() } @@ } + +impl PartialEq for AttrsRef<'_> { + #[inline] + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +}Stronger still: make the representation private so
Emptyremains the only zero-length form.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/binary/src/node.rs` around lines 393 - 420, The enum AttrsRef currently exposes its variants (pub enum AttrsRef<'a>), letting callers create Slice(Box::new([])) and breaking empty semantics; fix this by making the representation private (e.g. change the enum visibility from pub enum AttrsRef<'a> to pub(crate) enum AttrsRef<'a> or remove pub) so external code cannot construct Slice directly, and keep public APIs (from_vec, len, is_empty) as the only way to create/inspect AttrsRef; alternatively, if you must keep the enum public, remove the derived PartialEq and implement a custom PartialEq and adjust is_empty to treat AttrsRef::Slice(s) with s.len()==0 the same as AttrsRef::Empty (ensure from_vec still converts empty Vec to Empty).
🤖 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/keepalive.rs`:
- Around line 36-37: Add a unit test in the keepalive.rs tests that constructs
or simulates an IqError::EncodeError and asserts it maps to
KeepaliveResult::FatalFailure (matching the existing match arm for
IqError::EncodeError in the keepalive classification logic). Locate the matching
logic in keepalive.rs (the match handling IqError::InternalChannelClosed and
IqError::EncodeError) and write a test that invokes the classification path (or
calls the function that converts IqError to KeepaliveResult) to ensure
EncodeError is covered and asserts equality with KeepaliveResult::FatalFailure.
In `@src/request.rs`:
- Around line 169-185: Pre-allocate the encode buffer to avoid reallocations in
the fast path: replace the zero-capacity Vec::new() used before calling
spec.encode_iq_direct(&req_id, &mut buf) with a Vec::with_capacity(...) sized
for typical prekey IQs (e.g. 4 * 1024) or, if available, use a size hint from
the spec (e.g. spec.encoded_len_hint() or similar) to set capacity; keep the
rest of the logic around send_and_wait_iq(req_id, ... async {
self.send_raw_bytes(buf).await }) and parse_response(...) unchanged and still
map encode errors to IqError::EncodeError.
In `@wacore/binary/src/node.rs`:
- Around line 449-453: The safety comment for AttrsRef needs to be rewritten to
match the current enum representation (Empty | Slice(Box<[T]>)) and to justify
the transmute correctly: explain that Empty carries no lifetime, Slice contains
Box<[T]> which is covariant in T (and therefore in 'a), and thus AttrsRef is
covariant in 'a; update the comment above the transmute and the
_assert_attrs_ref_covariant function to refer to these exact variants and the
Box<[T]> covariance rather than the removed One/Many(Vec<T>) variants so the
unsafe justification matches the actual layout and invariants.
---
Duplicate comments:
In `@wacore/binary/src/node.rs`:
- Around line 393-420: The enum AttrsRef currently exposes its variants (pub
enum AttrsRef<'a>), letting callers create Slice(Box::new([])) and breaking
empty semantics; fix this by making the representation private (e.g. change the
enum visibility from pub enum AttrsRef<'a> to pub(crate) enum AttrsRef<'a> or
remove pub) so external code cannot construct Slice directly, and keep public
APIs (from_vec, len, is_empty) as the only way to create/inspect AttrsRef;
alternatively, if you must keep the enum public, remove the derived PartialEq
and implement a custom PartialEq and adjust is_empty to treat AttrsRef::Slice(s)
with s.len()==0 the same as AttrsRef::Empty (ensure from_vec still converts
empty Vec to Empty).
🪄 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: 78860f4d-f168-4d65-a6df-f94078cd19a9
📒 Files selected for processing (4)
src/keepalive.rssrc/request.rswacore/binary/src/decoder.rswacore/binary/src/node.rs
| | IqError::InternalChannelClosed | ||
| | IqError::EncodeError(_) => KeepaliveResult::FatalFailure, |
There was a problem hiding this comment.
Classification is correct, but add a unit test for EncodeError.
The fatal classification makes sense—encoding failures indicate a fundamental issue (e.g., malformed spec implementation), not a recoverable network condition. The tests cover all other IqError variants but miss EncodeError.
🧪 Suggested test
+ #[test]
+ fn test_classify_encode_error_is_fatal() {
+ assert_eq!(
+ classify_keepalive_error(&IqError::EncodeError(anyhow::anyhow!("encoding failed"))),
+ KeepaliveResult::FatalFailure,
+ );
+ }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/keepalive.rs` around lines 36 - 37, Add a unit test in the keepalive.rs
tests that constructs or simulates an IqError::EncodeError and asserts it maps
to KeepaliveResult::FatalFailure (matching the existing match arm for
IqError::EncodeError in the keepalive classification logic). Locate the matching
logic in keepalive.rs (the match handling IqError::InternalChannelClosed and
IqError::EncodeError) and write a test that invokes the classification path (or
calls the function that converts IqError to KeepaliveResult) to ensure
EncodeError is covered and asserts equality with KeepaliveResult::FatalFailure.
| { | ||
| 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) => {} | ||
| } | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Consider pre-allocating the buffer for the fast path.
Vec::new() starts at zero capacity and grows during encoding. Since prekey IQs are typically 1-4KB, a pre-sized buffer could avoid reallocations.
♻️ Optional optimization
- let mut buf = Vec::new();
+ let mut buf = Vec::with_capacity(2048);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@src/request.rs` around lines 169 - 185, Pre-allocate the encode buffer to
avoid reallocations in the fast path: replace the zero-capacity Vec::new() used
before calling spec.encode_iq_direct(&req_id, &mut buf) with a
Vec::with_capacity(...) sized for typical prekey IQs (e.g. 4 * 1024) or, if
available, use a size hint from the spec (e.g. spec.encoded_len_hint() or
similar) to set capacity; keep the rest of the logic around
send_and_wait_iq(req_id, ... async { self.send_raw_bytes(buf).await }) and
parse_response(...) unchanged and still map encode errors to
IqError::EncodeError.
|
@claude review |
|
Claude finished @jlucaso1's task in 4m 6s —— View job Review of PR #552
I've read every changed file and cross-referenced with the existing codebase. The PR has gone through multiple review rounds and many issues have been addressed. Here's what remains. Overall AssessmentThis is solid performance work. The Remaining Issues1. No equivalence test for direct encode vs build_iq + marshal (High)
A byte-comparison test that runs both paths with the same spec and asserts identical output would make this safe to maintain. This was flagged in the previous review and in CodeRabbit's review but hasn't been added yet. 2.
|
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
src/keepalive.rs (1)
36-37:⚠️ Potential issue | 🟡 MinorClassification is correct; add unit test for coverage.
Treating
EncodeErroras fatal is appropriate—encoding failures indicate a fundamental implementation issue, not a transient network condition. However, the test suite covers all otherIqErrorvariants but missesEncodeError.🧪 Suggested test
+ #[test] + fn test_classify_encode_error_is_fatal() { + assert_eq!( + classify_keepalive_error(&IqError::EncodeError(anyhow::anyhow!("encoding failed"))), + KeepaliveResult::FatalFailure, + ); + }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/keepalive.rs` around lines 36 - 37, Add a unit test in src/keepalive.rs that exercises the match arm for IqError::EncodeError and asserts it maps to KeepaliveResult::FatalFailure: create a dummy IqError::EncodeError instance, call the function (or internal mapping) that contains the shown match (the code handling IqError::InternalChannelClosed / IqError::EncodeError(_)), and assert the result equals KeepaliveResult::FatalFailure; if the mapping helper is private, place the test in the same module so it can access internals.wacore/binary/src/node.rs (1)
393-421:⚠️ Potential issue | 🟠 MajorHide or normalize the public
AttrsRefrepresentation.
AttrsRef::Slice(Box::new([]))is still constructible outside this module, so a logically empty value can reportfalsefromis_empty()and compare unequal toAttrsRef::EmptybecausePartialEqis derived on the variant shape. Either make the representation private, or makeis_empty()/PartialEqcanonicalize throughas_slice().Minimal mitigation if you keep the enum public
-#[derive(Debug, Clone, PartialEq)] +#[derive(Debug, Clone)] pub enum AttrsRef<'a> { Empty, Slice(Box<[(NodeStr<'a>, ValueRef<'a>)]>), } @@ #[inline] pub fn is_empty(&self) -> bool { - matches!(self, Self::Empty) + self.as_slice().is_empty() } } + +impl PartialEq for AttrsRef<'_> { + fn eq(&self, other: &Self) -> bool { + self.as_slice() == other.as_slice() + } +}Run this to see whether
AttrsRefvariants are constructed or pattern-matched directly outside this module before deciding between a private representation and the minimal mitigation:#!/bin/bash rg -n -C2 --type=rust '\bAttrsRef::(Empty|Slice)\b'🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@wacore/binary/src/node.rs` around lines 393 - 421, AttrsRef currently exposes its variants so external code can construct AttrsRef::Slice(Box::new([])) which breaks logical emptiness and PartialEq; fix by either making the representation private (remove public visibility of the enum variants so only this module can construct variants and keep pub AttrsRef type but not its variants) or by canonicalizing comparisons and emptiness: add an as_slice(&self) -> &[(NodeStr<'a>, ValueRef<'a>)] helper and change is_empty(&self) to check as_slice().is_empty(), and replace the derived PartialEq with a manual impl PartialEq for AttrsRef that compares via as_slice() (also ensure from_vec still constructs Slice/Empty appropriately). Ensure references to AttrsRef::Slice/Empty from outside are prevented if you choose the visibility change.
🤖 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/prekeys.rs`:
- Around line 426-441: The code currently calls
encoder.write_list_start(spec.pre_keys.len()) when encoding pre_keys; add a unit
test that constructs a Spec (or the struct used in this diff) with spec.pre_keys
= vec![] and asserts the encoded output equals build_iq's output for an empty
pre_keys case to ensure write_list_start(0) produces the expected binary; locate
the encoding path around write_list_start, write_string("list"),
write_string("key"), and write_string("id")/write_string("value") in the
functions handling prekey encoding (the block using spec.pre_keys) and add the
empty-list test exercising that exact branch.
---
Duplicate comments:
In `@src/keepalive.rs`:
- Around line 36-37: Add a unit test in src/keepalive.rs that exercises the
match arm for IqError::EncodeError and asserts it maps to
KeepaliveResult::FatalFailure: create a dummy IqError::EncodeError instance,
call the function (or internal mapping) that contains the shown match (the code
handling IqError::InternalChannelClosed / IqError::EncodeError(_)), and assert
the result equals KeepaliveResult::FatalFailure; if the mapping helper is
private, place the test in the same module so it can access internals.
In `@wacore/binary/src/node.rs`:
- Around line 393-421: AttrsRef currently exposes its variants so external code
can construct AttrsRef::Slice(Box::new([])) which breaks logical emptiness and
PartialEq; fix by either making the representation private (remove public
visibility of the enum variants so only this module can construct variants and
keep pub AttrsRef type but not its variants) or by canonicalizing comparisons
and emptiness: add an as_slice(&self) -> &[(NodeStr<'a>, ValueRef<'a>)] helper
and change is_empty(&self) to check as_slice().is_empty(), and replace the
derived PartialEq with a manual impl PartialEq for AttrsRef that compares via
as_slice() (also ensure from_vec still constructs Slice/Empty appropriately).
Ensure references to AttrsRef::Slice/Empty from outside are prevented if you
choose the visibility change.
🪄 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: ece36d00-1cd1-49b7-b751-9eaa34b29b99
📒 Files selected for processing (6)
.gitignoresrc/keepalive.rssrc/request.rswacore/binary/src/decoder.rswacore/binary/src/node.rswacore/src/iq/prekeys.rs
Derived PartialEq treated Empty != Slice(Box::new([])) even though both are logically empty. Manual PartialEq compares via as_slice() so all representations of the same content are equal. is_empty() now delegates to as_slice().is_empty() for consistency.
Verify direct encode matches build_iq + marshal for empty prekey list (write_list_start(0)) and single prekey. Extracted shared helper to avoid test duplication.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/binary/src/node.rs`:
- Around line 388-389: The comment in node.rs that states "16 bytes" and "24
bytes" for Box<[T]>/Vec<T> is architecture-specific; update the comment near the
`Box<[T]>`/`Vec<T>` explanation (and any mention of `NodeRef`) to either prefix
those numeric sizes with "on 64-bit targets" or replace them with a relative
description (e.g., "smaller footprint than Vec<T> on typical 64-bit platforms")
so the rationale remains correct across architectures.
🪄 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: 138f3b29-95ea-4c79-a036-db850d9d8c89
📒 Files selected for processing (1)
wacore/binary/src/node.rs
| /// Uses `Box<[T]>` (16 bytes: ptr + len) instead of `Vec<T>` (24 bytes: ptr + len + cap) | ||
| /// or inline storage (which inflated NodeRef size). Zero-attr nodes skip allocation |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Qualify the size numbers as 64-bit-specific.
16 bytes and 24 bytes are only true on 64-bit targets. Either qualify this as “on 64-bit” or reword it as a relative size comparison so the rationale stays accurate across architectures.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@wacore/binary/src/node.rs` around lines 388 - 389, The comment in node.rs
that states "16 bytes" and "24 bytes" for Box<[T]>/Vec<T> is
architecture-specific; update the comment near the `Box<[T]>`/`Vec<T>`
explanation (and any mention of `NodeRef`) to either prefix those numeric sizes
with "on 64-bit targets" or replace them with a relative description (e.g.,
"smaller footprint than Vec<T> on typical 64-bit platforms") so the rationale
remains correct across architectures.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9e268ef861
ℹ️ 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".
| // 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(); |
There was a problem hiding this comment.
Preserve send buffer capacity across frame handoff
Using out_buf.split().freeze() here does not keep the reusable buffer capacity as intended: split() is split_to(len), which advances out_buf to the tail region. Under sustained traffic, each send reduces the remaining capacity by that frame size until it reaches zero, so encode_frame_into starts reallocating frequently. This turns the new hot path into allocation churn on long-lived connections and undermines the performance goal of this change.
Useful? React with 👍 / 👎.
Summary
Three allocation optimizations targeting remaining hotspots from DHAT profiling, building on PR #551.
Combined impact (both PRs, DHAT full session):
This PR alone:
Changes
1. Zero-copy frame send via BytesMut
encode_frame_intois now generic over aFrameBuftrait (works with bothVec<u8>andBytesMut)BytesMutfor framing outputsplit().freeze()yields aBytesview without copying -- theBytesMutretains its allocationBytes::copy_from_sliceper outbound message2. Direct binary encoding for prekey upload IQ
Encoder,EncodeNode,ByteWriterand key write methods public inwacore-binaryencode_iq_directmethod toIqSpectrait (default returnsfalse)PreKeyUploadSpecimplements fast path with customEncodeNode-- writes the entire<iq>stanza directly via Encoder, bypassing NodeBuilder + marshal3. Covariant
AttrsRefenumtype AttrsRef<'a> = Vec<...>withEmpty | Slice(Box<[T]>)enumBox<[T]>is 16 bytes (ptr + len) vs Vec's 24 bytes -- no size regressionEmptyvariantyoke::Yokeable), compile-time assertion guards thisYokeableimpl with safety documentation4. DRY + correctness fixes
send_and_wait_iqshared helper (eliminates duplicatedsend_iq_raw)IqError::EncodeErrorvariant for direct-encode failures (was incorrectlyParseError)execute()reuses pre-generatedreq_idon fallback pathEncoder::new_vecclears buffer before writing format bytewrite_list_startrejects len > u16::MAX instead of silently truncatingencode_iq_directtrait contractTest plan
cargo test --all(949 tests pass)cargo clippy --all --testsclean (0 warnings)Summary by CodeRabbit
New Features
Performance
Bug Fixes
Tests