io: add WriterTo to MultiReader

Third version, previous: https://golang.org/cl/388455

This patch allows to zerocopy using MultiReader.
This is done by MultiReader implementing WriterTo.

Each sub reader is copied using usual io copy helper and thus use
WriterTo or ReadFrom with reflection.

There is a special case for when a subreader is a MultiReader.
Instead of using copyBuffer which would call multiReader.WriteTo,
multiReader.writeToWithBuffer is used instead, the difference
is that the temporary copy buffer is passed along, saving
allocations for nested MultiReaders.

The workflow looks like this:
- multiReader.WriteTo (allocates 32k buffer)
  - multiReader.writeToWithBuffer
    - for each subReader:
      - is instance of multiReader ?
        - yes, call multiReader.writeToWithBuffer
        - no, call copyBuffer(writer, currentReader, buffer)
          - does currentReader implements WriterTo ?
           - yes, use use currentReader.WriteTo
           - no, does writer implement ReadFrom ?
             - yes, use writer.ReadFrom
             - no, copy using Read / Write with buffer

This can be improved by lazy allocating the 32k buffer.
For example a MultiReader of such types:
  MultiReader(
    bytes.Reader, // WriterTo-able
    bytes.Reader, // WriterTo-able
    bytes.Reader, // WriterTo-able
  )

Doesn't need any allocation, all copy can be done using bytes.Reader's
internal data slice. However currently we still allocate a 32k buffer
for nothing.

This optimisation has been omited for a future patch because of high
complexity costs for a non obvious performance cost (it need a benchmark).
This patch at least is on par with the previous multiReader.Read
workflow allocation wise.

Fixes #50842
This commit is contained in:
Jorropo 2022-03-05 17:46:16 +01:00
parent 06a43e4ab6
commit 8ebe60ceac
2 changed files with 50 additions and 0 deletions

View File

@ -41,6 +41,31 @@ func (mr *multiReader) Read(p []byte) (n int, err error) {
return 0, EOF
}
func (mr *multiReader) WriteTo(w Writer) (sum int64, err error) {
return mr.writeToWithBuffer(w, make([]byte, 1024 * 32))
}
func (mr *multiReader) writeToWithBuffer(w Writer, buf []byte) (sum int64, err error) {
for i, r := range mr.readers {
var n int64
if subMr, ok := r.(*multiReader); ok { // reuse buffer with nested multiReaders
n, err = subMr.writeToWithBuffer(w, buf)
} else {
n, err = copyBuffer(w, r, buf)
}
sum += n
if err != nil {
mr.readers = mr.readers[i:] // permit resume / retry after error
return sum, err
}
mr.readers[i] = nil // permit early GC
}
mr.readers = nil
return sum, nil
}
var _ WriterTo = (*multiReader)(nil)
// MultiReader returns a Reader that's the logical concatenation of
// the provided input readers. They're read sequentially. Once all
// inputs have returned EOF, Read will return EOF. If any of the readers

View File

@ -63,6 +63,31 @@ func TestMultiReader(t *testing.T) {
})
}
func TestMultiReaderAsWriterTo(t *testing.T) {
mr := MultiReader(
strings.NewReader("foo "),
MultiReader( // Tickle the buffer reusing codepath
strings.NewReader(""),
strings.NewReader("bar"),
),
)
mrAsWriterTo, ok := mr.(WriterTo)
if !ok {
t.Fatalf("expected cast to WriterTo to succeed")
}
sink := &strings.Builder{}
n, err := mrAsWriterTo.WriteTo(sink)
if err != nil {
t.Fatalf("expected no error; got %v", err)
}
if n != 7 {
t.Errorf("expected read 7 bytes; got %d", n)
}
if result := sink.String(); result != "foo bar" {
t.Errorf(`expected "foo bar"; got %q`, result)
}
}
func TestMultiWriter(t *testing.T) {
sink := new(bytes.Buffer)
// Hide bytes.Buffer's WriteString method: