7871: Implement vectored writes for write_buf

[martin-augment]

7871: To review by AI

[gemini-code-assist]

Summary of Changes

Hello @martin-augment, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces an optimization to Tokio's write_buf utility by integrating support for vectored I/O. By conditionally utilizing poll_write_vectored when the underlying writer supports it, this change aims to enhance the efficiency of buffer writing operations, particularly for scenarios involving fragmented data, without altering the existing API for users.

Highlights

• Vectored Write Implementation: The write_buf utility now supports vectored writes, allowing it to efficiently write data from multiple non-contiguous buffers if the underlying AsyncWrite implementation supports it.
• Conditional Dispatch: The poll method in WriteBuf was updated to check is_write_vectored() on the underlying writer and dispatch to poll_write_vectored() when available, otherwise falling back to poll_write().
• New Test Case: A dedicated test write_buf_vectored has been added to ensure the correct behavior and functionality of the new vectored write implementation.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[coderabbitai]

Walkthrough

This pull request adds vectorized write support to the write_buf utility. The implementation branches based on whether the writer supports vectored operations via is_write_vectored. When true, it uses chunks_vectored to prepare up to 64 IoSlice entries and calls poll_write_vectored; otherwise it uses the existing poll_write path. A new test validates the vectorized write pathway with a custom AsyncWrite implementation that consumes input in 4-byte chunks via vectored writes.

✨ Finishing touches

• 📝 Generate docstrings

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[claude]

Code Review: Implement vectored writes for write_buf

Overview

This PR adds vectored write support to write_buf, which can improve performance when writing buffers that consist of multiple non-contiguous chunks. The implementation follows the existing pattern established in write_all_buf and is consistent with the codebase.

---

✅ Strengths

1. Consistent Pattern: The implementation matches the pattern used in write_all_buf.rs:44-55, maintaining consistency across the codebase.

2. Proper Constant Choice: Using MAX_VECTOR_ELEMENTS: usize = 64 is consistent with:

   • write_all_buf.rs (line 45)
   • tokio-util/src/util/poll_buf.rs (line 126, uses MAX_BUFS: usize = 64)

3. Good Test Coverage: The test properly validates:

   • Vectored writes are actually used (via panic! in poll_write)
   • The correct amount of data is written
   • The buffer position advances correctly
   • Multiple buffers are handled via poll_write_vectored

4. Correct Pin Usage: The pinning of mutable references (Pin::new(&mut *me.writer)) is done correctly in both branches.

---

🔍 Potential Issues

1. Zero-Write Handling (Minor but Important)

Unlike write_all_buf.rs:57-59, this implementation doesn't check for zero-byte writes. While write_buf only attempts a single write (not a loop like write_all_buf), it's still possible for poll_write_vectored to return Ok(0) without an error.

Current behavior in write_buf.rs:53-62:

let n = if me.writer.is_write_vectored() {
    let mut slices = [IoSlice::new(&[]); MAX_VECTOR_ELEMENTS];
    let cnt = me.buf.chunks_vectored(&mut slices);
    ready!(Pin::new(&mut *me.writer).poll_write_vectored(cx, &slices[..cnt]))?
} else {
    ready!(Pin::new(&mut *me.writer).poll_write(cx, me.buf.chunk()))?
};

me.buf.advance(n);
Poll::Ready(Ok(n))

If n == 0 but me.buf.has_remaining() is true, this could indicate a writer issue. The current code returns Ok(0), which may be acceptable for write_buf (since it's a partial write), but it's worth considering whether this should be documented or handled explicitly.

Suggestion: Consider adding a comment explaining that zero-byte writes are acceptable for write_buf (unlike write_all_buf), or verify this is the intended behavior.

2. Edge Case: Empty Buffer with chunks_vectored

If chunks_vectored returns slices that are all empty (cnt > 0 but all slices have len 0), the behavior depends on the writer's poll_write_vectored implementation. The check at line 49 (if !me.buf.has_remaining()) should prevent this, but it's worth noting.

---

📊 Performance Considerations

Positive: Vectored writes can significantly improve performance for:

• Buffers backed by rope-like data structures
• Chained buffers (common in the bytes crate)
• Network I/O where syscall overhead dominates

Stack Usage: The [IoSlice::new(&[]); MAX_VECTOR_ELEMENTS] array allocates 64 × sizeof(IoSlice) on the stack. On 64-bit systems, this is typically 64 × 16 = 1024 bytes, which is reasonable but worth noting for embedded contexts.

---

🔒 Security Considerations

No security concerns identified. The implementation:

• Properly bounds the slice access with &slices[..cnt]
• Uses safe Rust throughout
• Correctly advances the buffer by the exact amount written

---

🧪 Test Coverage Assessment

Good:

• Tests vectored write path explicitly
• Verifies non-vectored path isn't called
• Validates data integrity and buffer advancement

Could Add (optional enhancements):

1. Test with buffer that has more than 64 chunks (to verify MAX_VECTOR_ELEMENTS limit)
2. Test with empty buffer (Cursor::new(&[])), though line 49 handles this
3. Test with partially written vectored buffers (if the writer only writes some slices)

These are nice-to-haves rather than requirements, as the current test validates the core functionality.

---

📝 Minor Notes

1. Import Consistency: The import change from std::io to std::io::{self, IoSlice} (line 6) is clean and follows Rust conventions.

2. Code Duplication: There's some duplication with write_all_buf.rs and poll_write_buf in tokio-util, but this is acceptable given that:

   • Each function has slightly different semantics
   • The duplication is minimal
   • Extracting to a helper might add unnecessary complexity

---

✅ Recommendation

Approve with minor suggestion: The implementation is solid and follows established patterns. The only consideration is documenting or explicitly handling the zero-write case, though this may already be acceptable behavior for write_buf.

The code is:

• ✅ High quality and consistent with the codebase
• ✅ No bugs identified
• ✅ Good performance characteristics
• ✅ Secure
• ✅ Well tested for the core functionality

Great work on this PR! 🎉

[gemini-code-assist]

Code Review

This pull request introduces an optimization to write_buf by leveraging vectored writes when the underlying AsyncWrite implementation supports them. The changes look solid, including the addition of a new test case to validate the vectored write path. I have one minor suggestion to improve code style.

[gemini-code-assist]

For better readability and to follow the common convention in this crate, consider defining this constant at the module level rather than inside the poll method. This helps in grouping constants and makes the function body cleaner.

[martin-augment]

value:valid-but-wont-fix; category:bug; feedback: The Gemini AI reviewer is not correct! The constant is specific to the poll() method only and this is the reason why it is defined there. It should not be used (even by mistake) at other places in the code.

[augmentcode]

🤖 Augment PR Summary

Summary: This PR updates Tokio’s write_buf helper to take advantage of vectored writes when the underlying AsyncWrite supports them.

Changes:

• Switched `WriteBuf::poll` to build an `IoSlice` array via `Buf::chunks_vectored` and call `poll_write_vectored` when `is_write_vectored()` is true
• Kept the existing `poll_write` fallback path for non-vectored writers
• Capped the gathered slices to a fixed maximum (64) to bound stack usage
• Added a new test ensuring `write_buf` routes through `poll_write_vectored` for vectored-capable writers

Technical Notes: The buffer is still advanced by the returned byte count, matching the semantics used by other Tokio write utilities (e.g., write_all_buf).

_{🤖 Was this summary useful? React with 👍 or 👎}

[augmentcode]

Review completed. No suggestions at this time.

Comment augment review to trigger a new review at any time.