io: add WriterTo to MultiReader

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 omitted for a future patch because of high
complexity costs for a non obvious performance cost (it needs a benchmark).
This patch at least is on par with the previous MultiReader.Read
workflow allocation wise.

Fixes #50842

Change-Id: Ib070c8f36337d9dd86090df8a703c5df97a773ae
GitHub-Last-Rev: 8ebe60ceacec6bd52b63d9bdc05cd5b4ada57a6e
GitHub-Pull-Request: golang/go#51502
Reviewed-on: https://go-review.googlesource.com/c/go/+/390215
Run-TryBot: Ian Lance Taylor <iant@golang.org>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Ian Lance Taylor <iant@golang.org>
Reviewed-by: Daniel Martí <mvdan@mvdan.cc>
Trust: Daniel Martí <mvdan@mvdan.cc>

diff --git a/src/io/multi.go b/src/io/multi.go
index 24ee71e4ca65b19c19b83c366e5fcdb48571ac4f..909b7e45233a91adfb6f1576458298851d2596c3 100644 (file)
--- a/src/io/multi.go
+++ b/src/io/multi.go
@@ -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

diff --git a/src/io/multi_test.go b/src/io/multi_test.go
index e877e54571c4a31d2cbeeb1c308d93f79b0b5a7e..679312c23b01c518d1fe6bb896fc610a1cb5fc1c 100644 (file)
--- a/src/io/multi_test.go
+++ b/src/io/multi_test.go
@@ -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: