Efficient implementation of vectored writes for BufWriter

[mzabaluev]

Motivation

With vectored writes reintroduced in #3149, BufWriter can provide an efficient implementation both with or without vectored write support in the underlying writer, and advertise it with is_write_vectored.

Solution

Provide a custom implementation of AsyncWrite::poll_write_vectored by coalescing data from the slices in the buffer, unless it is more efficient to pass one or all of them directly to the underlying writer. This needs to be optimized differently depending on whether the underlying writer supports efficient vectored output.

[carllerche]

Thanks for the submission. Could we add some tests for the new functions?

[mzabaluev]

I had some tests written for #3092, but these relied on an async-friendly write_vectored method I implemented there, so I've been waiting for an equivalent to be reintroduced in tokio or tokio-util. #3156 only added poll_write_buf. I can make do with a temporary helper function, but would prefer to use AsyncWriteExt::write_buf once it's updated to use vectored output.

[mzabaluev]

I have rebased the branch and added some tests.

[Darksonn]

I'm sorry for taking so long to review this. The PR got lost among all the other PRs we get.

[Darksonn]

In general, I think that this could be simplified quite a lot if you make use of the fact that it is possible to reassign bufs to a sub-slice of itself to remove slices from the beginning of the buffer. E.g. to remove the first slice after having used it:

bufs = &bufs[1..];

[mzabaluev]

Thank you for the suggestion. I could not simplify by much, but hopefully the code is more readable now without an explicit iterator.

[Darksonn]

I gave it a try myself and came up with this. What do you think?

} else {
    // Remove empty buffers at the beginning of bufs.
    while bufs.first().map(|buf| buf.len()) == Some(0) {
        bufs = &bufs[1..];
    }
    if bufs.is_empty() {
        return Poll::Ready(Ok(0));
    }

    // Flush if the first buffer doesn't fit.
    let first_len = bufs[0].len();
    if first_len >= self.buf.capacity() - self.buf.len() {
        ready!(self.as_mut().flush_buf(cx))?;
        debug_assert!(self.buf.is_empty());
    }

    let me = self.project();

    if first_len >= me.buf.capacity() {
        // The slice is at least as large as the buffering capacity,
        // so it's better to write it directly, bypassing the buffer.
        debug_assert!(me.buf.is_empty());
        return me.inner.poll_write(cx, &*bufs[0]);
    }

    // Append the buffers that fit in the internal buffer.
    // The first buffer is guaranteed to fit due to the flush above.
    let mut total_written = 0;
    for buf in bufs {
        if buf.len() > me.buf.capacity() - me.buf.len() {
            break;
        } else {
            me.buf.extend_from_slice(&*buf);
            total_written += buf.len();
        }
    }

    debug_assert!(total_written != 0);
    Poll::Ready(Ok(total_written))
}

[mzabaluev]

@Darksonn Thanks! I have unrolled the first iteration of appending though, since we're into improving performance and the check in the loop body would be done redundantly, and there is a suitable else branch for that.

[Darksonn]

This seems good now, thanks.